Upload app.py with huggingface_hub

app.py

@@ -2,208 +2,198 @@ import gradio as gr
 import yaml
 import re
 import os
-import json
 
+# Paths from installed glmocr package
 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 ────────────────────────────────────────
 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")
+
+# ── STEP 2: Fix result_formatter.py ──────────────────────────
 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'"]:
+
+labels_to_remove = [
+    '"header"', "'header'",
+    '"footer"', "'footer'",
+    '"doc_header"', "'doc_header'",
+    '"doc_footer"', "'doc_footer'"
+]
+for label in labels_to_remove:
     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)
 
-ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
+print("✅ result_formatter.py fixed")
 
+ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
 
+# ── STEP 3: Extract header/footer from PDF text layer ─────────
+# The MaaS API only returns 'text' and 'table' labels — it never
+# labels anything as 'header' or 'footer'. So we extract those
+# directly from the PDF text layer using PyMuPDF instead.
+def extract_pdf_headers(path):
+    import fitz
+    page_headers = []
+    doc = fitz.open(path)
+    for page in doc:
+        page_rect   = page.rect
+        page_height = page_rect.height
+        # Top 12% = header area
+        header_rect = fitz.Rect(0, 0, page_rect.width, page_height * 0.12)
+        # Bottom 12% = footer area
+        footer_rect = fitz.Rect(0, page_height * 0.88, page_rect.width, page_height)
+        header_text = page.get_text(clip=header_rect).strip()
+        footer_text = page.get_text(clip=footer_rect).strip()
+        page_headers.append({"header": header_text, "footer": footer_text})
+    doc.close()
+    return page_headers
+
+# ── STEP 4: Main OCR function ─────────────────────────────────
 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: API key not set.", "Failed."
+
     try:
         from glmocr import parse
+
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
+
+        # Convert PDF to images + extract headers from text layer
         if path.lower().endswith(".pdf"):
             import fitz
+            pdf_headers = extract_pdf_headers(path)
             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)
             doc.close()
             result = parse(page_images)
         else:
+            pdf_headers = []
             result = parse(path)
 
-        debug_lines = ["========== DEBUG =========="]
-        debug_lines.append(f"1. result type: {type(result).__name__}")
-        debug_lines.append(f"2. result length: {len(result) if isinstance(result, list) else 'N/A'}")
-        if isinstance(result, list) and len(result) > 0:
-            first = result[0]
-            debug_lines.append(f"3. first element type: {type(first).__name__}")
-            debug_lines.append(f"4. has markdown_result: {hasattr(first, 'markdown_result')}")
-            debug_lines.append(f"5. has json_result: {hasattr(first, 'json_result')}")
-            if hasattr(first, "markdown_result"):
-                md = getattr(first, "markdown_result", None)
-                debug_lines.append(f"6. markdown_result len: {len(md) if md else 0}")
-            if hasattr(first, "json_result"):
-                j = getattr(first, "json_result", None)
-                debug_lines.append(f"7. json_result len: {len(j) if j else 0}")
-        debug_lines.append("==========================")
-
-        r0 = result[0] if (isinstance(result, list) and len(result) > 0) else None
-        raw_resp = getattr(r0, '_maas_response', None) if r0 else None
-        if raw_resp is not None:
-            try:
-                resp_str = json.dumps(raw_resp, indent=2, ensure_ascii=False)
-            except Exception:
-                resp_str = str(raw_resp)
-            if len(resp_str) > 15000:
-                resp_str = resp_str[:15000] + "\n\n... (truncated, total " + str(len(resp_str)) + " chars)"
-            debug_lines.append("10. FULL API RESPONSE (first page):")
-            debug_lines.append(resp_str)
-        else:
-            debug_lines.append("10. FULL API RESPONSE: not available (no _maas_response)")
-
-        if not isinstance(result, list):
-            result = [result]
-        if not result:
-            return "(No content)", "\n".join(debug_lines)
-
-        first = result[0]
-        has_markdown = getattr(first, "markdown_result", None) is not None
-        has_json = getattr(first, "json_result", None) is not None
-
-        if has_markdown or has_json:
-            markdown_parts = []
-            pages_data = []
-            for r in result:
-                md = getattr(r, "markdown_result", None) or ""
-                markdown_parts.append(md.strip() if md else "(empty page)")
-                j = getattr(r, "json_result", None) or []
-                for page in j:
-                    pages_data.append(page)
-            markdown = "\n\n---\n\n".join(markdown_parts)
-
-            labels_seen = set()
-            for page in pages_data:
-                if not isinstance(page, list):
-                    continue
-                for r in page:
-                    if isinstance(r, dict):
-                        labels_seen.add(r.get("label", ""))
-            debug_lines.append(f"8. ALL labels from API: {sorted(labels_seen)}")
-            debug_lines.append(f"9. markdown_parts: {len(markdown_parts)}, pages_data: {len(pages_data)}")
-
-            total_headers = 0
-            total_footers = 0
-            all_summary = []
-            for page_num, page_regions in enumerate(pages_data):
-                if not isinstance(page_regions, list):
-                    page_regions = [page_regions] if page_regions else []
-                page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]
-                for region in page_regions:
-                    if not isinstance(region, dict):
-                        continue
-                    label = region.get("label", region.get("type", "text"))
-                    content = str(region.get("content", region.get("text", "")))
-                    if label in ABANDON:
-                        continue
-                    if label == "header":
-                        total_headers += 1
-                        page_summary.append(f"🔵 HEADER: {content[:100]}")
-                    elif label == "footer":
-                        total_footers += 1
-                        page_summary.append(f"🟢 FOOTER: {content[:100]}")
-                    else:
-                        page_summary.append(f"[{label}]: {content[:100]}")
-                all_summary.extend(page_summary)
-                all_summary.append("")
-            summary_body = (
-                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)
-            )
-            summary = "\n".join(debug_lines) + "\n\n" + summary_body
-            return markdown, summary
-
-        pages_data = result if isinstance(result, list) else []
+        # Parse result
+        pages_data = result.json_result if hasattr(result, "json_result") else []
+        if not isinstance(pages_data, list):
+            pages_data = [pages_data] if pages_data else []
+
         total_headers = 0
         total_footers = 0
-        all_markdown = []
-        all_summary = []
+        all_markdown  = []
+        all_summary   = []
+
         for page_num, page_regions in enumerate(pages_data):
             if not isinstance(page_regions, list):
                 page_regions = [page_regions] if page_regions else []
-            page_md = []
+
+            page_md      = []
             page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]
+
+            # Inject header/footer from PDF text layer
+            # (API never returns these as separate labeled regions)
+            if pdf_headers and page_num < len(pdf_headers):
+                hdr = pdf_headers[page_num]["header"]
+                ftr = pdf_headers[page_num]["footer"]
+                if hdr:
+                    total_headers += 1
+                    page_md.append("<!-- HEADER -->\n" + hdr)
+                    page_summary.append("🔵 HEADER: " + hdr[:100])
+                if ftr:
+                    total_footers += 1
+                    page_md.append("<!-- FOOTER -->\n" + ftr)
+                    page_summary.append("🟢 FOOTER: " + ftr[:100])
+
+            # Process API regions (text, table, etc.)
             for region in page_regions:
                 if not isinstance(region, dict):
                     continue
-                label = region.get("label", "text")
+                label   = region.get("label", "text")
                 content = str(region.get("content", ""))
+
                 if label in ABANDON:
                     continue
+
+                # API returns 'header'/'footer' labels in some docs
+                # Keep them too in case the API does label them
                 if label == "header":
                     total_headers += 1
-                    page_summary.append(f"🔵 HEADER: {content[:100]}")
-                    page_md.append(f"<!-- HEADER -->\n{content}")
+                    page_summary.append("🔵 HEADER (API): " + content[:100])
+                    page_md.append("<!-- HEADER -->\n" + content)
                 elif label == "footer":
                     total_footers += 1
-                    page_summary.append(f"🟢 FOOTER: {content[:100]}")
-                    page_md.append(f"<!-- FOOTER -->\n{content}")
+                    page_summary.append("🟢 FOOTER (API): " + content[:100])
+                    page_md.append("<!-- FOOTER -->\n" + content)
                 else:
-                    page_summary.append(f"[{label}]: {content[:100]}")
+                    page_summary.append("[" + label + "]: " + content[:100])
                     page_md.append(content)
+
             all_markdown.append("\n\n".join(page_md))
             all_summary.extend(page_summary)
             all_summary.append("")
-        summary_body = (
-            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)
+
+        summary = (
+            "Total pages   : " + str(len(pages_data)) + "\n"
+            "Headers found : " + str(total_headers) + "\n"
+            "Footers found : " + str(total_footers) + "\n"
+            + "─" * 40 + "\n"
+            + "\n".join(all_summary)
         )
+
         markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
-        summary = "\n".join(debug_lines) + "\n\n" + summary_body
         return markdown, summary
 
     except Exception as e:
         import traceback
-        return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."
+        return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
+
+# ── STEP 5: 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 🟢 now appear in output.
+    Multi-page PDFs fully supported.
+    """)
+
+    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.")
-    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")
             markdown_out = gr.Textbox(lines=30, label="")
         with gr.Column():
-            gr.Markdown("### 🗂️ Detected Regions (+ Debug)")
+            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])
+
+    run_btn.click(fn=run_ocr, inputs=file_input,
+                  outputs=[markdown_out, regions_out])
+
 demo.launch()