Update app.py

app.py

@@ -1,127 +1,54 @@
-import gradio as gr
-import yaml
-import re
 import os
-import tempfile
+import re
+import yaml
+import gradio as gr
 
 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_HEADER_FRAC = 0.12
-FALLBACK_FOOTER_FRAC = 0.12
-OCR_FALLBACK_HEADER_FOOTER = True  # When text layer is empty, OCR the crop
+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: Remove header/footer from abandon lists only ──────
-with open(config_path, "r") as f:
+# ─── 1. Config: include headers and footers (remove from abandon lists) ───
+with open(CONFIG_PATH, "r") as f:
     config = yaml.safe_load(f)
 
 config["pipeline"]["maas"]["enabled"] = True
-config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
+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
 
 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]
-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]
-
-with open(config_path, "w") as f:
+for key in ("layout", "result_formatter"):
+    if key not in config.get("pipeline", {}):
+        continue
+    section = config["pipeline"][key]
+    if key == "layout" and "label_task_mapping" in section:
+        abandon = section["label_task_mapping"].get("abandon")
+        if isinstance(abandon, list):
+            section["label_task_mapping"]["abandon"] = [x for x in abandon if x not in to_remove]
+    elif key == "result_formatter" and "abandon" in section:
+        abandon = section["abandon"]
+        if isinstance(abandon, list):
+            section["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")
-
-# ── STEP 2: Fix result_formatter.py ──────────────────────────
-with open(formatter_path, "r") as f:
+# ─── 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'"]:
-    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:
+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):
-    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, None
-    return min(y_tops) / norm_height, max(y_bottoms) / norm_height
-
 
-def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
-    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
-
-
-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):
-    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
-    md = ""
-    if hasattr(page_result, "markdown_result") and page_result.markdown_result:
-        md = page_result.markdown_result.strip()
-    return md, regions
-
-
-# ── STEP 4: Main OCR ─────────────────────────────────────────
+# ─── 3. Run OCR (no custom header/footer logic; pipeline returns full content) ───
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
-
     try:
         from glmocr import parse
         import pymupdf as fitz
@@ -139,63 +66,29 @@ def run_ocr(uploaded_file):
                 page_images.append(img_path)
             doc.close()
             results = parse(page_images)
-            if not isinstance(results, list):
-                results = [results]
         else:
-            page_images = []
             results = parse(path)
-            if not isinstance(results, list):
-                results = [results]
-
-        all_pages_md = []
-
-        for page_num, page_result in enumerate(results):
-            page_md, regions = get_page_data(page_result)
-            parts = []
 
-            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)
+        if not isinstance(results, list):
+            results = [results]
 
-                # 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:
-                if page_md:
-                    parts.append(page_md)
-
-            if parts:
-                all_pages_md.append("\n\n".join(parts))
-
-        final = "\n\n---\n\n".join(all_pages_md) if all_pages_md else "(No content)"
+        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
-
     except Exception as e:
         import traceback
-        return "Error: " + str(e) + "\n\n" + traceback.format_exc()
-
+        return f"Error: {e}\n\n{traceback.format_exc()}"
 
-# ── STEP 5: Gradio UI ─────────────────────────────────────────
+# ─── 4. UI ───
 with gr.Blocks(title="GLM-OCR") as demo:
-    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)
+    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers and footers are 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()