Update app.py

app.py

@@ -4,8 +4,8 @@ import re
 import os
 
 import glmocr
-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")
 
 # ── STEP 1: Fix config ────────────────────────────────────────
@@ -14,7 +14,6 @@ with open(config_path, "r") as f:
 
 config["pipeline"]["maas"]["enabled"] = True
 config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
-
 config["pipeline"]["result_formatter"]["abandon"] = [
     "number", "footnote", "aside_text",
     "reference", "footer_image", "header_image",
@@ -34,7 +33,7 @@ labels_to_remove = [
     '"header"', "'header'",
     '"footer"', "'footer'",
     '"doc_header"', "'doc_header'",
-    '"doc_footer"', "'doc_footer'"
+    '"doc_footer"', "'doc_footer'",
 ]
 for label in labels_to_remove:
     source = re.sub(r',\s*' + re.escape(label), '', source)
@@ -46,37 +45,42 @@ with open(formatter_path, "w") as f:
 
 print("✅ result_formatter.py fixed")
 
-# ── STEP 3: Helper functions ──────────────────────────────────
-def get_top_gap(regions, img_height):
+# ── STEP 3: Helpers ───────────────────────────────────────────
+def get_header_footer_zones(regions, norm_height=1000):
     """
-    Find the top gap the API missed using bbox_2d coordinates.
-    Only extracts header (top gap) — never bottom — to avoid
-    duplicating body content the API already captured correctly.
+    From API regions' bbox_2d (normalized 0 to norm_height), return
+    (header_end_frac, footer_start_frac). No hardcoded percentages.
     """
+    if not regions:
+        return None, None
     y_tops = []
+    y_bottoms = []
     for r in regions:
         bbox = r.get("bbox_2d")
         if bbox and len(bbox) == 4:
             y_tops.append(bbox[1])
+            y_bottoms.append(bbox[3])
     if not y_tops:
-        return None
-    first_y = min(y_tops)
-    return (first_y / img_height) if first_y > img_height * 0.08 else None
+        return None, None
+    header_end = min(y_tops) / norm_height
+    footer_start = max(y_bottoms) / norm_height
+    return header_end, footer_start
 
 
 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)
+    doc = fitz.open(pdf_path)
     page = doc[page_num]
     h, w = page.rect.height, page.rect.width
-    text = page.get_text(clip=fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)).strip()
+    rect = fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)
+    text = page.get_text(clip=rect).strip()
     doc.close()
     return text
 
 
 def get_page_data(page_result):
-    """Extract markdown and raw regions from one PipelineResult."""
+    """Return (markdown, regions) for one PipelineResult."""
     regions = []
     if hasattr(page_result, "json_result"):
         jr = page_result.json_result
@@ -90,7 +94,7 @@ def get_page_data(page_result):
     return md, regions
 
 
-# ── STEP 4: Main OCR function ─────────────────────────────────
+# ── STEP 4: Main OCR ─────────────────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
@@ -99,26 +103,22 @@ def run_ocr(uploaded_file):
         from glmocr import parse
         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")
 
         if is_pdf:
-            doc          = fitz.open(path)
-            page_images  = []
-            page_heights = []
+            doc = fitz.open(path)
+            page_images = []
             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 = f"/tmp/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 = parse(page_images)
             if not isinstance(results, list):
                 results = [results]
         else:
-            page_heights = []
             results = parse(path)
             if not isinstance(results, list):
                 results = [results]
@@ -129,21 +129,21 @@ def run_ocr(uploaded_file):
             page_md, regions = get_page_data(page_result)
             parts = []
 
-            if is_pdf and page_num < len(page_heights):
-                top_gap = get_top_gap(regions, page_heights[page_num])
+            if is_pdf:
+                header_end, footer_start = get_header_footer_zones(regions, 1000)
 
-                # HEADER at TOP — from PDF text layer (API misses this zone)
-                if top_gap is not None:
-                    hdr = extract_zone_text(path, page_num, 0, top_gap)
-                    if hdr:
-                        parts.append(hdr)
+                if header_end is not None and header_end > 0:
+                    hdr = extract_zone_text(path, page_num, 0, header_end)
+                    if hdr and hdr.strip():
+                        parts.append(hdr.strip())
 
-                # BODY — from API (includes all content + footer text correctly)
                 if page_md:
                     parts.append(page_md)
 
-                # No bottom zone extraction — API already captures footer content.
-                # Extracting bottom zone causes duplication of body content.
+                if footer_start is not None and footer_start < 1.0:
+                    ftr = extract_zone_text(path, page_num, footer_start, 1.0)
+                    if ftr and ftr.strip():
+                        parts.append(ftr.strip())
             else:
                 if page_md:
                     parts.append(page_md)
@@ -151,7 +151,6 @@ def run_ocr(uploaded_file):
             if parts:
                 all_pages_md.append("\n\n".join(parts))
 
-        # Exactly like original SDK output format
         final = "\n\n---\n\n".join(all_pages_md) if all_pages_md else "(No content)"
         return final
 
@@ -164,17 +163,16 @@ def run_ocr(uploaded_file):
 with gr.Blocks(title="GLM-OCR") as demo:
     gr.Markdown("""
     # 🔍 GLM-OCR
-    Upload a PDF or image to extract text.
-    Headers included at top of each page.
+    Upload a PDF or image. Headers and footers from the PDF text layer (all PDFs).
     """)
 
     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")
+    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()
+demo.launch()