Update app.py

app.py

@@ -1,8 +1,8 @@
 """
-GLM-OCR Hugging Face Space app with client-side header/footer fallback for MaaS.
-Works for every PDF and every image: uses API bboxes when available, else minimal band.
+GLM-OCR Hugging Face Space — MaaS mode with header extraction.
+Header is extracted from PDF text layer using bbox gap detection.
+Footer is NOT extracted separately — API already captures it as text regions.
 """
-# Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
 import asyncio
 try:
     _orig_close = asyncio.BaseEventLoop.close
@@ -21,84 +21,44 @@ import re
 import tempfile
 import yaml
 import gradio as gr
-
 import glmocr
 
 log = logging.getLogger("glmocr_app")
 
-GLMOCR_BASE = os.path.dirname(glmocr.__file__)
-CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
+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")
 
-# When MaaS is enabled, the cloud API ignores local layout config and does not return
-# header/footer regions. We always run client-side OCR on top/bottom bands as fallback.
-DEFAULT_ZONE_FRAC = float(os.environ.get("GLMOCR_DEFAULT_ZONE_FRAC", "0.12"))
-# MaaS rejects very small images (400). Only send crops that meet minimum dimensions.
-MIN_CROP_HEIGHT = int(os.environ.get("GLMOCR_MIN_CROP_HEIGHT", "112"))
-MIN_CROP_PIXELS = int(os.environ.get("GLMOCR_MIN_CROP_PIXELS", "12544"))  # 112*112
-# Wider PDF bands for position-based extraction (capture account numbers etc. in top/bottom)
-PDF_HEADER_BAND_FRAC = float(os.environ.get("GLMOCR_PDF_HEADER_BAND", "0.15"))   # top 15%
-PDF_FOOTER_BAND_FRAC = float(os.environ.get("GLMOCR_PDF_FOOTER_BAND", "0.85"))   # bottom 15%
-
-# Single shared parser to avoid "GLM-OCR initialized" per request and asyncio cleanup issues.
-_parser = None
-
-def get_parser():
-    global _parser
-    if _parser is None:
-        from glmocr import GlmOcr
-        _parser = GlmOcr(
-            api_key=os.environ.get("GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"),
-            mode="maas",
-        )
-    return _parser
+DEFAULT_HEADER_FRAC = float(os.environ.get("GLMOCR_DEFAULT_HEADER_FRAC", "0.12"))
 
-# ---------------------------------------------------------------------------
-# 1. Config: set MaaS and optionally include headers/footers for non-MaaS
-# ---------------------------------------------------------------------------
+# ── STEP 1: Fix config ────────────────────────────────────────
 try:
     with open(CONFIG_PATH, "r") as f:
         config = yaml.safe_load(f)
+
     config["pipeline"]["maas"]["enabled"] = True
-    config["pipeline"]["maas"]["api_key"] = os.environ.get("GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV")
-
-    to_include_as_text = {"header", "footer"}
-    layout_section = config.get("pipeline", {}).get("layout", {})
-    if "label_task_mapping" in layout_section:
-        mapping = layout_section["label_task_mapping"]
-        abandon = mapping.get("abandon")
-        if isinstance(abandon, list):
-            mapping["abandon"] = [x for x in abandon if x not in to_include_as_text]
-        text_labels = mapping.get("text")
-        if isinstance(text_labels, list):
-            for label in to_include_as_text:
-                if label not in text_labels:
-                    text_labels.append(label)
+    config["pipeline"]["maas"]["api_key"] = os.environ.get(
+        "GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
+    )
 
+    to_include = {"header", "footer"}
     formatter_section = config.get("pipeline", {}).get("result_formatter", {})
-    if "abandon" in formatter_section:
-        abandon = formatter_section["abandon"]
-        if isinstance(abandon, list):
-            formatter_section["abandon"] = [x for x in abandon if x not in to_include_as_text]
-    if "label_visualization_mapping" in formatter_section:
-        text_vis = formatter_section["label_visualization_mapping"].get("text")
-        if isinstance(text_vis, list):
-            for label in to_include_as_text:
-                if label not in text_vis:
-                    text_vis.append(label)
+    if "abandon" in formatter_section and isinstance(formatter_section["abandon"], list):
+        formatter_section["abandon"] = [
+            x for x in formatter_section["abandon"] if x not in to_include
+        ]
 
     with open(CONFIG_PATH, "w") as f:
         yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 except Exception:
-    pass  # e.g. read-only package on Hugging Face; MaaS ignores layout anyway
+    pass
 
-# ---------------------------------------------------------------------------
-# 2. result_formatter: stop stripping header/footer (best-effort; may be read-only)
-# ---------------------------------------------------------------------------
+# ── STEP 2: Fix result_formatter.py ──────────────────────────
 try:
     with open(FORMATTER_PATH, "r") as f:
         source = f.read()
-    for label in ('"header"', "'header'", '"footer"', "'footer'", '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"):
+    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)
@@ -107,211 +67,144 @@ try:
 except Exception:
     pass
 
+# ── Single shared parser ──────────────────────────────────────
+_parser = None
 
-def get_header_footer_zones(regions, norm_height=1000):
-    if not regions:
-        return None, None
-    y_tops, y_bottoms = [], []
+def get_parser():
+    global _parser
+    if _parser is None:
+        from glmocr import GlmOcr
+        _parser = GlmOcr(
+            api_key=os.environ.get("GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"),
+            mode="maas",
+        )
+    return _parser
+
+# ── STEP 3: Header extraction helpers ────────────────────────
+def get_top_gap_frac(regions, img_height):
+    """
+    Find the fraction of image height that the API missed at the top.
+    Uses bbox_2d of the first (topmost) region.
+    Returns a fraction (0–1) or None if no gap detected.
+    """
+    y_tops = []
     for r in regions:
         bbox = r.get("bbox_2d") if isinstance(r, dict) else getattr(r, "bbox_2d", None)
         if bbox and len(bbox) >= 4:
             y_tops.append(bbox[1])
-            y_bottoms.append(bbox[3])
     if not y_tops:
-        return None, None
-    return min(y_tops) / norm_height, max(y_bottoms) / norm_height
+        return DEFAULT_HEADER_FRAC  # no bbox info — use default band
+    first_y = min(y_tops)
+    # Only treat as missed header if gap > 8% of image height
+    return (first_y / img_height) if first_y > img_height * 0.08 else None
 
 
-def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract text from a horizontal band using a clip rect."""
+def extract_header_text(pdf_path, page_num, y_end_frac):
+    """Extract text from the top zone of a PDF page using PyMuPDF."""
     try:
         import pymupdf as fitz
-        doc = fitz.open(pdf_path)
+        doc  = fitz.open(pdf_path)
         page = doc[page_num]
         h, w = page.rect.height, page.rect.width
-        rect = fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)
-        text = page.get_text(clip=rect).strip()
+        text = page.get_text(clip=fitz.Rect(0, 0, w, h * y_end_frac)).strip()
         doc.close()
         return text
     except Exception:
         return ""
 
 
-def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract all text whose word bbox falls in the given vertical band (by position).
-    Use a wider band (e.g. top 15% / bottom 15%) to capture header/footer that clip might miss."""
-    try:
-        import pymupdf as fitz
-        doc = fitz.open(pdf_path)
-        page = doc[page_num]
-        h = page.rect.height
-        y_lo = h * y_start_frac
-        y_hi = h * y_end_frac
-        words = page.get_text("words")  # list of (x0, y0, x1, y1, word, ...)
-        doc.close()
-        parts = []
-        for w in words:
-            if len(w) >= 5:
-                y0, y1 = float(w[1]), float(w[3])
-                if y0 < y_hi and y1 > y_lo:
-                    parts.append(w[4])
-        return " ".join(parts).strip()
-    except Exception:
-        return ""
-
-
-def ocr_zone(image_path, y_start_frac, y_end_frac):
-    """Run OCR on a horizontal band. Pads small crops to meet API minimum size instead of skipping."""
-    zone_name = "header" if y_end_frac < 0.5 else "footer"
-    try:
-        from PIL import Image
-        img = Image.open(image_path).convert("RGB")
-        w, h = img.size
-        y0 = max(0, int(h * y_start_frac))
-        y1 = min(h, int(h * y_end_frac))
-        if y1 <= y0:
-            return ""
-        crop = img.crop((0, y0, w, y1))
-        cw, ch = crop.size
-        # Pad to minimum size so MaaS accepts the image (avoids 400 on tiny strips)
-        if ch < MIN_CROP_HEIGHT or (cw * ch) < MIN_CROP_PIXELS:
-            need_h = max(ch, MIN_CROP_HEIGHT)
-            need_w = max(cw, 1)
-            if (need_w * need_h) < MIN_CROP_PIXELS:
-                need_w = max(need_w, (MIN_CROP_PIXELS + need_h - 1) // need_h)
-            canvas = Image.new("RGB", (need_w, need_h), (255, 255, 255))
-            if zone_name == "header":
-                canvas.paste(crop, (0, 0))  # crop at top
-            else:
-                canvas.paste(crop, (0, need_h - ch))  # crop at bottom
-            crop = canvas
-            cw, ch = crop.size
-        fd, path = tempfile.mkstemp(suffix=".jpg")
-        os.close(fd)
-        try:
-            crop.save(path, "JPEG", quality=92)
-            parser = get_parser()
-            out = parser.parse(path)
-            if not isinstance(out, list):
-                out = [out]
-            if out and getattr(out[0], "markdown_result", None):
-                text = (out[0].markdown_result or "").strip()
-                return text
-        finally:
-            try:
-                os.unlink(path)
-            except Exception:
-                pass
-    except Exception as e:
-        log.warning("[%s] ocr_zone failed: %s", zone_name, e, exc_info=True)
-    return ""
-
-
-def get_page_md_and_regions(page_result):
+def get_page_data(page_result):
+    """Extract markdown and regions from one PipelineResult."""
     md = ""
     if hasattr(page_result, "markdown_result") and page_result.markdown_result:
         md = (page_result.markdown_result or "").strip()
     regions = []
     if hasattr(page_result, "json_result"):
         jr = page_result.json_result
-        if isinstance(jr, dict) and "regions" in jr:
-            regions = jr.get("regions") or []
-        elif isinstance(jr, list) and len(jr) > 0:
+        if isinstance(jr, list) and len(jr) > 0:
             r = jr[0] if isinstance(jr[0], list) else jr
             if isinstance(r, list):
                 regions = r
-            elif isinstance(r, dict) and "regions" in r:
-                regions = r.get("regions") or []
     return md, regions
 
 
+# ── STEP 4: Main OCR function ─────────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
     try:
         import pymupdf as fitz
 
-        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
+        path   = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
         is_pdf = path.lower().endswith(".pdf")
         parser = get_parser()
 
         if is_pdf:
-            doc = fitz.open(path)
-            page_images = []
+            doc          = fitz.open(path)
+            page_images  = []
+            page_heights = []
             for i in range(len(doc)):
-                pix = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
+                pix      = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
                 img_path = os.path.join(tempfile.gettempdir(), f"maas_page_{i}.png")
                 pix.save(img_path)
                 page_images.append(img_path)
+                # Image height at 1.5x matches bbox_2d coordinate space
+                page_heights.append(doc[i].rect.height * 1.5)
             doc.close()
             results = parser.parse(page_images)
         else:
-            page_images = [path]
-            results = parser.parse(path)
+            page_images  = [path]
+            page_heights = []
+            results      = parser.parse(path)
 
         if not isinstance(results, list):
             results = [results]
 
         all_pages = []
-        for page_num, page_result in enumerate(results):
-            page_md, regions = get_page_md_and_regions(page_result)
-            header_end_frac, footer_start_frac = get_header_footer_zones(regions, 1000)
-
-            # MaaS does not return header/footer regions; always use at least default bands
-            he = header_end_frac if header_end_frac is not None else DEFAULT_ZONE_FRAC
-            fs = footer_start_frac if footer_start_frac is not None else (1.0 - DEFAULT_ZONE_FRAC)
-            if he <= 0:
-                he = DEFAULT_ZONE_FRAC  # ensure we always OCR a top band
-            if fs >= 1.0:
-                fs = 1.0 - DEFAULT_ZONE_FRAC  # ensure we always OCR a bottom band
 
+        for page_num, page_result in enumerate(results):
+            page_md, regions = get_page_data(page_result)
             parts = []
 
-            # Always run header band (top 8–12% of page)
-            if page_num < len(page_images):
-                img_path = page_images[page_num]
-                hdr = ""
-                if is_pdf:
-                    hdr = extract_zone_text_pdf(path, page_num, 0, he)
-                    if not (hdr and hdr.strip()):
-                        hdr = extract_pdf_text_in_band(path, page_num, 0, PDF_HEADER_BAND_FRAC)
-                if not (hdr and hdr.strip()):
-                    hdr = ocr_zone(img_path, 0, he)
-                if hdr and hdr.strip():
-                    parts.append(hdr.strip())
-
+            # ── HEADER at TOP ─────────────────────────────────
+            # Only for PDFs — extract top zone the API missed.
+            # We do NOT extract footer — the API already captures
+            # footer text (copyright, address) as regular text regions.
+            if is_pdf and page_num < len(page_heights):
+                top_gap = get_top_gap_frac(regions, page_heights[page_num])
+                if top_gap is not None:
+                    hdr = extract_header_text(path, page_num, top_gap)
+                    if hdr:
+                        parts.append(hdr)
+
+            # ── BODY in MIDDLE ────────────────────────────────
             if page_md:
                 parts.append(page_md)
 
-            # Always run footer band (bottom 8–12% of page)
-            if page_num < len(page_images):
-                img_path = page_images[page_num]
-                ftr = ""
-                if is_pdf:
-                    ftr = extract_zone_text_pdf(path, page_num, fs, 1.0)
-                    if not (ftr and ftr.strip()):
-                        ftr = extract_pdf_text_in_band(path, page_num, PDF_FOOTER_BAND_FRAC, 1.0)
-                if not (ftr and ftr.strip()):
-                    ftr = ocr_zone(img_path, fs, 1.0)
-                if ftr and ftr.strip():
-                    parts.append(ftr.strip())
+            # No footer extraction — avoids duplicating body content.
+            # The API captures real footer text as text regions in page_md.
 
             if parts:
                 all_pages.append("\n\n".join(parts))
 
-        return "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)"
+        return "\n\n---\n\n".join(all_pages) if all_pages else "(No content)"
+
     except Exception as e:
         import traceback
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
 
+# ── STEP 5: Gradio UI ─────────────────────────────────────────
 with gr.Blocks(title="GLM-OCR") as demo:
-    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers and footers included.")
-    file_in = gr.File(label="Upload PDF or image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
-    run_btn = gr.Button("Run OCR", variant="primary")
-    out = gr.Textbox(lines=40, label="Output")
+    gr.Markdown("# 🔍 GLM-OCR\nUpload a PDF or image. Headers included in correct position.")
+    file_in = gr.File(
+        label="Upload PDF or image",
+        file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"]
+    )
+    run_btn = gr.Button("▶  Run OCR", variant="primary")
+    out     = gr.Textbox(lines=40, label="Output")
     run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()