Download checkpoint from HF hub in OcrReorderPipeline

app.py

@@ -1,14 +1,11 @@
-import os
-import json
-import base64
+import os, json, base64
 from io import BytesIO
 from PIL import Image
 import gradio as gr
-
 from inference import OcrReorderPipeline
 from transformers import AutoProcessor, LayoutLMv3Model, AutoTokenizer
 
-# 1) Load your model + tokenizer + processor as before
+# Load model/tokenizer/processor...
 repo      = "Uddipan107/ocr-layoutlmv3-base-t5-small"
 model     = LayoutLMv3Model.from_pretrained(repo)
 tokenizer = AutoTokenizer.from_pretrained(repo, subfolder="preprocessor")
@@ -16,14 +13,17 @@ processor = AutoProcessor.from_pretrained(repo, subfolder="preprocessor", apply_
 pipe      = OcrReorderPipeline(model, tokenizer, processor, device=0)
 
 def infer(image_path, json_file):
-    # 2) Extract the filename user uploaded
     img_name = os.path.basename(image_path)
 
-    # 3) Load the entire JSON; assume it’s a list of entries
+    # Parse NDJSON from the uploaded file
+    data = []
     with open(json_file.name, "r", encoding="utf-8") as f:
-        data = json.load(f)
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            data.append(json.loads(line))
 
-    # 4) Find the entry matching this image
     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,24 +31,22 @@ def infer(image_path, json_file):
     words = entry["src_word_list"]
     boxes = entry["src_wordbox_list"]
 
-    # 5) Read the image, encode to base64 for your pipeline
+    # Read and encode image
     img = Image.open(image_path).convert("RGB")
     buf = BytesIO(); img.save(buf, format="PNG")
     b64 = base64.b64encode(buf.getvalue()).decode()
 
-    # 6) Call your pipeline and return the reordered text
+    # Run pipeline
     return pipe(b64, words, boxes)[0]
 
 demo = gr.Interface(
     fn=infer,
     inputs=[
-      # get the file path so we can match the filename
       gr.Image(type="filepath", label="Upload Image"),
-      # this is the JSON file containing a list of entries
-      gr.File(label="Upload JSON file")
+      gr.File(label="Upload JSON (NDJSON format)")
     ],
     outputs="text",
-    title="OCR Reorder (match image → JSON entry)"
+    title="OCR Reorder (Image + NDJSON upload)"
 )
 
 if __name__ == "__main__":