Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,136 @@
+
+import gradio as gr
+import yaml
+import re
+import os
+import torch
+
+# Paths from installed glmocr package (required on HF Space)
+import glmocr
+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 — 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"] = [
+    "number", "footnote", "aside_text", "reference",
+    "footer_image", "header_image",
+]
+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 (header & footer kept in output)")
+
+# ── STEP 2: Fix result_formatter.py (remove hardcoded header/footer) ───────
+with open(formatter_path, "r") as f:
+    source = f.read()
+labels_to_remove = [
+    '"header"', "'header'", '"footer"', "'footer'",
+    '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"
+]
+for label in labels_to_remove:
+    source = re.sub(r',\s*' + re.escape(label), '', source)
+    source = re.sub(re.escape(label) + r'\s*,', '', source)
+    source = re.sub(re.escape(label), '', source)
+with open(formatter_path, "w") as f:
+    f.write(source)
+print("✅ result_formatter.py fixed")
+
+# ── STEP 3: Load model ────────────────────────────────────────────────────
+from transformers import AutoProcessor, GlmOcrForConditionalGeneration
+print("Loading model... (~2GB first run)")
+processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
+model = GlmOcrForConditionalGeneration.from_pretrained(
+    "zai-org/GLM-OCR",
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+print("✅ Model ready on", next(model.parameters()).device)
+ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
+
+# ── 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},
+                {"type": "text", "text": "Document Parsing:"}
+            ]}
+        ]
+        inputs = processor.apply_chat_template(
+            messages, tokenize=True, add_generation_prompt=True,
+            return_dict=True, return_tensors="pt"
+        ).to(model.device)
+        inputs.pop("token_type_ids", None)
+        with torch.no_grad():
+            output_ids = model.generate(**inputs, max_new_tokens=2048)
+        raw = processor.decode(
+            output_ids[0][inputs["input_ids"].shape[1]:],
+            skip_special_tokens=False
+        )
+        raw = raw.replace("<|user|>", "").strip()
+        json_match = re.search(r'\[.*\]', raw, re.DOTALL)
+        regions = json.loads(json_match.group()) if json_match else []
+        header_count = footer_count = 0
+        region_lines = []
+        markdown_parts = []
+        for region in regions:
+            label = region.get("label", "text")
+            content = str(region.get("content", ""))
+            if label in ABANDON:
+                continue
+            if label == "header":
+                header_count += 1
+                region_lines.append("🔵 HEADER:\n" + content + "\n")
+                markdown_parts.append("<!-- HEADER -->\n" + content)
+            elif label == "footer":
+                footer_count += 1
+                region_lines.append("🟢 FOOTER:\n" + content + "\n")
+                markdown_parts.append("<!-- FOOTER -->\n" + content)
+            else:
+                region_lines.append("[" + label + "]: " + content[:150])
+                markdown_parts.append(content)
+        summary = (
+            "Headers found : " + str(header_count) + "\n"
+            "Footers found : " + str(footer_count) + "\n"
+            "Total regions : " + str(len(regions)) + "\n" + "─"*40 + "\n"
+            + "\n".join(region_lines)
+        )
+        markdown = "\n\n".join(markdown_parts) if markdown_parts else raw
+        return markdown, summary
+    except Exception as e:
+        import traceback
+        return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
+
+# ── 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():
+        with gr.Column():
+            gr.Markdown("### 📄 Markdown Output")
+            markdown_out = gr.Textbox(lines=25, label="")
+        with gr.Column():
+            gr.Markdown("### 🗂️ Detected Regions")
+            regions_out = gr.Textbox(lines=25, label="")
+    run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
+demo.launch()