Download checkpoint from HF hub in OcrReorderPipeline

app.py

@@ -1,21 +1,28 @@
-import os, json, base64
+import os
+import json
+import base64
 from io import BytesIO
 from PIL import Image
 import gradio as gr
+
 from inference import OcrReorderPipeline
 from transformers import AutoProcessor, LayoutLMv3Model, AutoTokenizer
 
-# Load model/tokenizer/processor...
+# ── 1) Load model/tokenizer/processor ─────────────────────────────────────
 repo      = "Uddipan107/ocr-layoutlmv3-base-t5-small"
 model     = LayoutLMv3Model.from_pretrained(repo)
 tokenizer = AutoTokenizer.from_pretrained(repo, subfolder="preprocessor")
 processor = AutoProcessor.from_pretrained(repo, subfolder="preprocessor", apply_ocr=False)
-pipe      = OcrReorderPipeline(model, tokenizer, processor, device=0)
 
+# instantiate your custom pipeline
+pipe = OcrReorderPipeline(model, tokenizer, processor, device=0)
+
+# ── 2) Inference function ─────────────────────────────────────────────────
 def infer(image_path, json_file):
+    # Extract filename
     img_name = os.path.basename(image_path)
 
-    # Parse NDJSON from the uploaded file
+    # Parse NDJSON (one JSON object per line)
     data = []
     with open(json_file.name, "r", encoding="utf-8") as f:
         for line in f:
@@ -24,6 +31,7 @@ def infer(image_path, json_file):
                 continue
             data.append(json.loads(line))
 
+    # Find the matching entry
     entry = next((e for e in data if e["img_name"] == img_name), None)
     if entry is None:
         return f"❌ No JSON entry found for image '{img_name}'"
@@ -31,22 +39,25 @@ def infer(image_path, json_file):
     words = entry["src_word_list"]
     boxes = entry["src_wordbox_list"]
 
-    # Read and encode image
+    # Load & encode the image as base64
     img = Image.open(image_path).convert("RGB")
-    buf = BytesIO(); img.save(buf, format="PNG")
+    buf = BytesIO()
+    img.save(buf, format="PNG")
     b64 = base64.b64encode(buf.getvalue()).decode()
 
-    # Run pipeline
-    return pipe(b64, words, boxes)[0]
+    # ▶️ Correct keyword call so preprocess gets all args
+    reordered = pipe(image=b64, words=words, boxes=boxes)[0]
+    return reordered
 
+# ── 3) Build the Gradio interface ───────────────────────────────────────────
 demo = gr.Interface(
     fn=infer,
     inputs=[
-      gr.Image(type="filepath", label="Upload Image"),
-      gr.File(label="Upload JSON (NDJSON format)")
+        gr.Image(type="filepath", label="Upload Image"),
+        gr.File(label="Upload JSON (NDJSON)")
     ],
     outputs="text",
-    title="OCR Reorder (Image + NDJSON upload)"
+    title="OCR Reorder Pipeline"
 )
 
 if __name__ == "__main__":