Update app.py

app.py

@@ -2,15 +2,16 @@ import gradio as gr
 import yaml
 import re
 import os
+import tempfile
 
 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")
 
-# Fallback band when API reports no gap (so all PDFs get header/footer)
-FALLBACK_HEADER_FRAC = 0.05
-FALLBACK_FOOTER_FRAC = 0.05
+FALLBACK_HEADER_FRAC = 0.12
+FALLBACK_FOOTER_FRAC = 0.12
+OCR_FALLBACK_HEADER_FOOTER = True  # When text layer is empty, OCR the crop
 
 # ── STEP 1: Remove header/footer from abandon lists only ──────
 with open(config_path, "r") as f:
@@ -20,53 +21,36 @@ config["pipeline"]["maas"]["enabled"] = True
 config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
 
 to_remove = {"header", "footer"}
-
 if "layout" in config["pipeline"] and "label_task_mapping" in config["pipeline"]["layout"]:
     abandon = config["pipeline"]["layout"]["label_task_mapping"].get("abandon")
     if isinstance(abandon, list):
-        config["pipeline"]["layout"]["label_task_mapping"]["abandon"] = [
-            x for x in abandon if x not in to_remove
-        ]
-
+        config["pipeline"]["layout"]["label_task_mapping"]["abandon"] = [x for x in abandon if x not in to_remove]
 if "result_formatter" in config["pipeline"]:
     abandon = config["pipeline"]["result_formatter"].get("abandon")
     if isinstance(abandon, list):
-        config["pipeline"]["result_formatter"]["abandon"] = [
-            x for x in abandon if x not in to_remove
-        ]
+        config["pipeline"]["result_formatter"]["abandon"] = [x for x in abandon if x not in to_remove]
 
 with open(config_path, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 
-print("✅ config.yaml updated (header/footer removed from abandon only)")
+print("✅ config.yaml updated")
 
 # ── STEP 2: Fix result_formatter.py ──────────────────────────
 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:
+for label in ['"header"', "'header'", '"footer"', "'footer'", '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"]:
     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: Helpers ───────────────────────────────────────────
 def get_header_footer_zones(regions, norm_height=1000):
-    """From bbox_2d (0–norm_height) return (header_end_frac, footer_start_frac)."""
     if not regions:
         return None, None
-    y_tops = []
-    y_bottoms = []
+    y_tops, y_bottoms = [], []
     for r in regions:
         bbox = r.get("bbox_2d")
         if bbox and len(bbox) == 4:
@@ -78,7 +62,6 @@ def get_header_footer_zones(regions, norm_height=1000):
 
 
 def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract text from a vertical slice of a PDF page."""
     import pymupdf as fitz
     doc = fitz.open(pdf_path)
     page = doc[page_num]
@@ -89,8 +72,38 @@ def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
     return text
 
 
+def ocr_image_zone(image_path, y_start_frac, y_end_frac, is_footer=False):
+    """Crop a band from the page image and run parse() on it; return markdown or ''."""
+    try:
+        from PIL import Image
+        from glmocr import parse
+        img = Image.open(image_path).convert("RGB")
+        w, h = img.size
+        y0 = int(h * y_start_frac)
+        y1 = int(h * y_end_frac)
+        if y1 <= y0:
+            return ""
+        crop = img.crop((0, y0, w, y1))
+        fd, path = tempfile.mkstemp(suffix=".png")
+        os.close(fd)
+        crop.save(path)
+        try:
+            out = parse(path)
+            if not isinstance(out, list):
+                out = [out]
+            if out and hasattr(out[0], "markdown_result") and out[0].markdown_result:
+                return out[0].markdown_result.strip()
+        finally:
+            try:
+                os.unlink(path)
+            except Exception:
+                pass
+    except Exception:
+        pass
+    return ""
+
+
 def get_page_data(page_result):
-    """Return (markdown, regions) for one PipelineResult."""
     regions = []
     if hasattr(page_result, "json_result"):
         jr = page_result.json_result
@@ -129,6 +142,7 @@ def run_ocr(uploaded_file):
             if not isinstance(results, list):
                 results = [results]
         else:
+            page_images = []
             results = parse(path)
             if not isinstance(results, list):
                 results = [results]
@@ -139,21 +153,27 @@ def run_ocr(uploaded_file):
             page_md, regions = get_page_data(page_result)
             parts = []
 
-            if is_pdf:
+            if is_pdf and page_num < len(page_images):
                 header_end, footer_start = get_header_footer_zones(regions, 1000)
                 he = max(header_end, FALLBACK_HEADER_FRAC) if header_end is not None else FALLBACK_HEADER_FRAC
                 fs = min(footer_start, 1.0 - FALLBACK_FOOTER_FRAC) if footer_start is not None else (1.0 - FALLBACK_FOOTER_FRAC)
 
+                # Header: PDF text layer first; if empty, OCR the top crop
                 if he > 0:
                     hdr = extract_zone_text(path, page_num, 0, he)
+                    if not (hdr and hdr.strip()) and OCR_FALLBACK_HEADER_FOOTER:
+                        hdr = ocr_image_zone(page_images[page_num], 0, FALLBACK_HEADER_FRAC, is_footer=False)
                     if hdr and hdr.strip():
                         parts.append(hdr.strip())
 
                 if page_md:
                     parts.append(page_md)
 
+                # Footer: PDF text layer first; if empty, OCR the bottom crop
                 if fs < 1.0:
                     ftr = extract_zone_text(path, page_num, fs, 1.0)
+                    if not (ftr and ftr.strip()) and OCR_FALLBACK_HEADER_FOOTER:
+                        ftr = ocr_image_zone(page_images[page_num], 1.0 - FALLBACK_FOOTER_FRAC, 1.0, is_footer=True)
                     if ftr and ftr.strip():
                         parts.append(ftr.strip())
             else:
@@ -173,18 +193,9 @@ def run_ocr(uploaded_file):
 
 # ── STEP 5: Gradio UI ─────────────────────────────────────────
 with gr.Blocks(title="GLM-OCR") as demo:
-    gr.Markdown("""
-    # 🔍 GLM-OCR
-    Upload a PDF or image. Headers and footers from PDF text layer (all PDFs).
-    """)
-
-    file_input = gr.File(
-        label="Upload PDF or Image",
-        file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"]
-    )
+    gr.Markdown("# 🔍 GLM-OCR\nUpload a PDF or image. Headers/footers from text layer or OCR fallback.")
+    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")
     output_box = gr.Textbox(lines=40, label="Output")
-
     run_btn.click(fn=run_ocr, inputs=file_input, outputs=output_box)
-
 demo.launch()