Update app.py

app.py

@@ -1,5 +1,10 @@
+"""
+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.
+"""
 import os
 import re
+import tempfile
 import yaml
 import gradio as gr
 
@@ -9,14 +14,16 @@ 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")
 
-# ─── 1. Config: include headers and footers (remove from abandon lists) ───
+DEFAULT_ZONE_FRAC = float(os.environ.get("GLMOCR_DEFAULT_ZONE_FRAC", "0.08"))
+
+# ---------------------------------------------------------------------------
+# 1. Config: include headers/footers (remove from abandon lists)
+# ---------------------------------------------------------------------------
 with open(CONFIG_PATH, "r") as f:
     config = yaml.safe_load(f)
 
 config["pipeline"]["maas"]["enabled"] = True
-api_key = os.environ.get("GLMOCR_MaaS_API_KEY") or os.environ.get("ZHIPU_API_KEY")
-if api_key:
-    config["pipeline"]["maas"]["api_key"] = api_key
+config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
 
 to_remove = {"header", "footer"}
 for key in ("layout", "result_formatter"):
@@ -35,7 +42,9 @@ for key in ("layout", "result_formatter"):
 with open(CONFIG_PATH, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 
-# ─── 2. result_formatter: stop stripping header/footer in postprocess ───
+# ---------------------------------------------------------------------------
+# 2. result_formatter: stop stripping header/footer in postprocess
+# ---------------------------------------------------------------------------
 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'"):
@@ -45,7 +54,81 @@ for label in ('"header"', "'header'", '"footer"', "'footer'", '"doc_header"', "'
 with open(FORMATTER_PATH, "w") as f:
     f.write(source)
 
-# ─── 3. Run OCR (no custom header/footer logic; pipeline returns full content) ───
+
+def get_header_footer_zones(regions, norm_height=1000):
+    if not regions:
+        return None, None
+    y_tops, y_bottoms = [], []
+    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
+
+
+def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
+    try:
+        import pymupdf as fitz
+        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()
+        doc.close()
+        return text
+    except Exception:
+        return ""
+
+
+def ocr_zone(image_path, y_start_frac, y_end_frac):
+    try:
+        from PIL import Image
+        from glmocr import parse
+        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))
+        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 getattr(out[0], "markdown_result", None):
+                return (out[0].markdown_result or "").strip()
+        finally:
+            try:
+                os.unlink(path)
+            except Exception:
+                pass
+    except Exception:
+        pass
+    return ""
+
+
+def get_page_md_and_regions(page_result):
+    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, 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
+
+
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
@@ -61,34 +144,70 @@ def run_ocr(uploaded_file):
             page_images = []
             for i in range(len(doc)):
                 pix = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
-                img_path = f"/tmp/maas_page_{i}.png"
+                img_path = os.path.join(tempfile.gettempdir(), f"maas_page_{i}.png")
                 pix.save(img_path)
                 page_images.append(img_path)
             doc.close()
             results = parse(page_images)
         else:
+            page_images = [path]
             results = parse(path)
 
         if not isinstance(results, list):
             results = [results]
 
-        parts = []
-        for r in results:
-            md = getattr(r, "markdown_result", None) or ""
-            if isinstance(md, str) and md.strip():
-                parts.append(md.strip())
-        final = "\n\n---\n\n".join(parts) if parts else "(No content)"
-        return final
+        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)
+
+            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 and header_end_frac is not None:
+                he = 0
+            if fs >= 1.0 and footer_start_frac is not None:
+                fs = 1.0
+
+            parts = []
+
+            if he > 0 and 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 = ocr_zone(img_path, 0, he)
+                if hdr and hdr.strip():
+                    parts.append(hdr.strip())
+
+            if page_md:
+                parts.append(page_md)
+
+            if fs < 1.0 and 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 = ocr_zone(img_path, fs, 1.0)
+                if ftr and ftr.strip():
+                    parts.append(ftr.strip())
+
+            if parts:
+                all_pages.append("\n\n".join(parts))
+
+        return "\n\n---\n\n".join(all_pages) if all_pages else "(No content)"
     except Exception as e:
         import traceback
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
-# ─── 4. UI ───
+
 with gr.Blocks(title="GLM-OCR") as demo:
-    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers and footers are included.")
+    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")
     run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()