Upload app.py with huggingface_hub

app.py

@@ -11,7 +11,7 @@ GLMOCR_BASE = os.path.dirname(glmocr.__file__)
 config_path = os.path.join(GLMOCR_BASE, "config.yaml")
 formatter_path = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 
-# ── STEP 1: Fix config.yaml on startup ───────────────────────
+# ── STEP 1: Fix config — keep header & footer (do NOT add them to abandon) ──
 with open(config_path, "r") as f:
     config = yaml.safe_load(f)
 config["pipeline"]["result_formatter"]["abandon"] = [
@@ -21,9 +21,9 @@ config["pipeline"]["result_formatter"]["abandon"] = [
 config["pipeline"]["enable_layout"] = True
 with open(config_path, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
-print("✅ config.yaml fixed")
+print("✅ config.yaml fixed (header & footer kept in output)")
 
-# ── STEP 2: Fix result_formatter.py on startup ───────────────
+# ── STEP 2: Fix result_formatter.py (remove hardcoded header/footer) ───────
 with open(formatter_path, "r") as f:
     source = f.read()
 labels_to_remove = [
@@ -38,7 +38,7 @@ with open(formatter_path, "w") as f:
     f.write(source)
 print("✅ result_formatter.py fixed")
 
-# ── STEP 3: Load model ────────────────────────────────────────
+# ── STEP 3: Load model ────────────────────────────────────────────────────
 from transformers import AutoProcessor, GlmOcrForConditionalGeneration
 print("Loading model... (~2GB first run)")
 processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
@@ -50,12 +50,25 @@ model = GlmOcrForConditionalGeneration.from_pretrained(
 print("✅ Model ready on", next(model.parameters()).device)
 ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
 
-# ── STEP 4: OCR function ──────────────────────────────────────
+# ── STEP 4: OCR (PDF → image then run model) ───────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file.", "No regions detected."
     try:
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
+        if path.lower().endswith(".pdf"):
+            try:
+                import fitz
+                doc = fitz.open(path)
+                page = doc[0]
+                pix = page.get_pixmap(matrix=fitz.Matrix(1, 1), alpha=False)
+                img_path = path[:-4] + "_page0.png"
+                pix.save(img_path)
+                doc.close()
+                path = img_path
+            except Exception as e:
+                return "PDF conversion failed: " + str(e), "Failed."
+
         messages = [
             {"role": "user", "content": [
                 {"type": "image", "url": path},
@@ -107,9 +120,9 @@ def run_ocr(uploaded_file):
         import traceback
         return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
 
-# ── STEP 5: Gradio UI ─────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
-    gr.Markdown("# 🔍 GLM-OCR — Header & Footer Fix\nUpload any PDF or image. Headers 🔵 and Footers 🟢 now appear in output.")
+# ── STEP 5: Gradio UI ──────────────────────────────────────────────────────
+with gr.Blocks(title="GLM-OCR — Header & Footer Kept") as demo:
+    gr.Markdown("# 🔍 GLM-OCR — Header & Footer Kept\nUpload PDF or image. Headers 🔵 and Footers 🟢 are kept in output.")
     file_input = gr.File(label="Upload PDF or Image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
     run_btn = gr.Button("▶  Run OCR", variant="primary", size="lg")
     with gr.Row():