Update app.py

app.py

@@ -3,87 +3,71 @@ import yaml
 import re
 import os
 
-# Paths from installed glmocr package
 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")
 
-# ── STEP 1: Load and fix config (MaaS + keep header/footer) ─────────────────
+# ── Config: MaaS on, keep header/footer ────────────────────────────────────
 with open(config_path, "r") as f:
     config = yaml.safe_load(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",
+    "number", "footnote", "aside_text", "reference", "footer_image", "header_image",
 ]
 config["pipeline"]["enable_layout"] = True
-
 with open(config_path, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 
-print("✅ config.yaml fixed (MaaS enabled, header/footer kept)")
-
-# ── STEP 2: Fix result_formatter.py ────────────────────────────────────────
+# ── Formatter: don't strip header/footer ───────────────────────────────────
 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")
-
 ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
 
-def _get_regions(page_item):
-    """Get list of regions from one page. Page can be list of regions or dict with 'regions' key."""
+def get_regions(page_item):
+    if page_item is None:
+        return []
     if isinstance(page_item, list):
         return page_item
     if isinstance(page_item, dict):
-        return page_item.get("regions", page_item.get("content", [page_item]))
+        for key in ("regions", "blocks", "items", "content_list"):
+            if key in page_item and isinstance(page_item[key], list):
+                return page_item[key]
+        if "content" in page_item or "text" in page_item:
+            return [page_item]
+    if isinstance(page_item, str):
+        return [{"label": "text", "content": page_item}]
     return []
 
-def _get_content(region):
-    """Get text from a region. Supports 'content', 'text', 'raw'."""
-    if not isinstance(region, dict):
-        return str(region)
-    return str(region.get("content") or region.get("text") or region.get("raw") or "")
-
-def _get_label(region):
-    """Get label from a region."""
-    if not isinstance(region, dict):
+def get_label(r):
+    if not isinstance(r, dict):
         return "text"
-    return region.get("label") or region.get("type") or "text"
+    return (r.get("label") or r.get("type") or r.get("block_type") or "text")
+
+def get_content(r):
+    if isinstance(r, str):
+        return r
+    if not isinstance(r, dict):
+        return str(r)
+    return str(r.get("content") or r.get("text") or r.get("raw") or "")
 
-# ── STEP 3: OCR via Zhipu MaaS (no local model) ─────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file.", "No regions detected."
-
     api_key = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
     if not api_key:
-        return (
-            "Error: ZHIPU_API_KEY is not set.",
-            "Failed."
-        )
-
+        return "Error: API key not set.", "Failed."
     try:
         from glmocr import parse
-
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
-
         if path.lower().endswith(".pdf"):
             import fitz
             doc = fitz.open(path)
@@ -97,28 +81,24 @@ def run_ocr(uploaded_file):
             result = parse(page_images)
         else:
             result = parse(path)
-
         pages_data = result if isinstance(result, list) else []
-
         total_headers = 0
         total_footers = 0
         all_markdown = []
         all_summary = []
-
         for page_num, page_item in enumerate(pages_data):
-            page_regions = _get_regions(page_item)
+            page_regions = get_regions(page_item)
             if not isinstance(page_regions, list):
-                page_regions = [page_regions] if page_regions else []
+                page_regions = [page_regions]
             page_md = []
             page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]
-
             for region in page_regions:
-                label = _get_label(region)
-                content = _get_content(region)
-
+                label = get_label(region)
+                content = get_content(region)
+                if not content.strip():
+                    continue
                 if label in ABANDON:
                     continue
-
                 if label == "header":
                     total_headers += 1
                     page_summary.append(f"🔵 HEADER: {content[:100]}")
@@ -130,37 +110,25 @@ def run_ocr(uploaded_file):
                 else:
                     page_summary.append(f"[{label}]: {content[:100]}")
                     page_md.append(content)
-
             all_markdown.append("\n\n".join(page_md))
             all_summary.extend(page_summary)
             all_summary.append("")
-
         summary = (
             f"Total pages   : {len(pages_data)}\n"
             f"Headers found : {total_headers}\n"
             f"Footers found : {total_footers}\n"
-            f"{'─'*40}\n"
-            + "\n".join(all_summary)
+            f"{'─'*40}\n" + "\n".join(all_summary)
         )
         markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
         return markdown, summary
-
     except Exception as e:
         import traceback
         return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."
 
-# ── STEP 4: Gradio UI ──────────────────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
-    gr.Markdown("""
-    # 🔍 GLM-OCR — Zhipu MaaS
-    Upload PDF or image. Headers 🔵 and Footers 🟢 from the **full pipeline** (layout + OCR).
-    Uses your Zhipu API key.
-    """)
-    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")
+with gr.Blocks(title="GLM-OCR — MaaS") as demo:
+    gr.Markdown("# 🔍 GLM-OCR — Zhipu MaaS\nUpload PDF or image. Headers 🔵 and Footers 🟢 when returned by the API.")
+    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")
     with gr.Row():
         with gr.Column():
             gr.Markdown("### 📄 Markdown Output")
@@ -169,5 +137,4 @@ with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
             gr.Markdown("### 🗂️ Detected Regions")
             regions_out = gr.Textbox(lines=30, label="")
     run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
-
 demo.launch()