Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import yaml
 import re
 import os
-import json
 
 import glmocr
 GLMOCR_BASE    = os.path.dirname(glmocr.__file__)
@@ -13,8 +12,8 @@ formatter_path = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 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"]["maas"]["enabled"] = True
+config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
 
 config["pipeline"]["result_formatter"]["abandon"] = [
     "number", "footnote", "aside_text",
@@ -47,33 +46,8 @@ with open(formatter_path, "w") as f:
 
 print("✅ result_formatter.py fixed")
 
-ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
-
-# ── STEP 3: Extract text from a specific zone of a PDF page ───
-# We use this ONLY to recover text from zones the API missed.
-# We figure out what zones were missed by looking at bbox_2d gaps.
-def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract text from a vertical fraction of a PDF page."""
-    import fitz
-    doc  = fitz.open(pdf_path)
-    page = doc[page_num]
-    h    = page.rect.height
-    w    = 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 get_missed_zones(regions, page_height, page_width, min_gap_frac=0.05):
-    """
-    Look at bbox_2d of API regions and find vertical gaps at the
-    top and bottom of the page that were not covered by any region.
-    Returns: (top_gap_frac, bottom_gap_frac) or None if no gap.
-    """
-    if not regions:
-        return None, None
-
-    # Collect all y-coordinates from bbox_2d
+# ── STEP 3: Helper functions ──────────────────────────────────
+def get_missed_zones(regions, img_height):
     y_tops    = []
     y_bottoms = []
     for r in regions:
@@ -81,274 +55,127 @@ def get_missed_zones(regions, page_height, page_width, min_gap_frac=0.05):
         if bbox and len(bbox) == 4:
             y_tops.append(bbox[1])
             y_bottoms.append(bbox[3])
-
     if not y_tops:
         return None, None
+    first_y = min(y_tops)
+    last_y  = max(y_bottoms)
+    top_gap    = (first_y / img_height) if first_y > img_height * 0.08 else None
+    bottom_gap = (last_y  / img_height) if last_y  < img_height * 0.92 else None
+    return top_gap, bottom_gap
 
-    first_region_y = min(y_tops)
-    last_region_y  = max(y_bottoms)
-
-    # Top gap: if first region starts below 10% of page height
-    top_gap = None
-    if first_region_y > page_height * 0.10:
-        top_gap = first_region_y / page_height   # as fraction
 
-    # Bottom gap: if last region ends above 90% of page height
-    bottom_gap = None
-    if last_region_y < page_height * 0.90:
-        bottom_gap = last_region_y / page_height  # as fraction
+def extract_zone_text(pdf_path, page_num, y_start_frac, y_end_frac):
+    import fitz
+    doc  = fitz.open(pdf_path)
+    page = doc[page_num]
+    h    = page.rect.height
+    w    = 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
 
-    return top_gap, bottom_gap
 
-# ── STEP 4: Parse one page result ────────────────────────────
-def get_page_regions(page_result):
+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:
-            regions = jr[0] if isinstance(jr[0], list) else jr
-            if isinstance(regions, list):
-                return regions, "json"
-    if hasattr(page_result, "markdown_result"):
-        md = page_result.markdown_result
-        if md and isinstance(md, str) and md.strip():
-            return md.strip(), "raw"
-    return [], "empty"
-
-# ── STEP 5: Build readable HTML from JSON ────────────────────
-def json_to_readable_html(pages_json):
-    html = """
-    <style>
-        .page-block { border: 2px solid #333; border-radius: 8px; margin-bottom: 24px; overflow: hidden; }
-        .page-title { background: #333; color: white; padding: 8px 16px; font-size: 16px; font-weight: bold; }
-        .region { padding: 10px 16px; border-bottom: 1px solid #eee; }
-        .region:last-child { border-bottom: none; }
-        .region-label { display: inline-block; padding: 2px 8px; border-radius: 4px; font-size: 12px; font-weight: bold; margin-bottom: 6px; }
-        .label-header { background: #3b82f6; color: white; }
-        .label-footer { background: #22c55e; color: white; }
-        .label-table  { background: #f59e0b; color: white; }
-        .label-text   { background: #6b7280; color: white; }
-        .label-other  { background: #8b5cf6; color: white; }
-        .region-content { font-size: 13px; color: #222; white-space: pre-wrap; word-break: break-word; }
-        .region-content table { border-collapse: collapse; width: 100%; margin-top: 4px; }
-        .region-content td, .region-content th { border: 1px solid #ccc; padding: 4px 8px; font-size: 12px; }
-        .region-content th { background: #f3f4f6; }
-        .source-badge { font-size: 10px; color: #999; margin-left: 8px; }
-        .empty-page { padding: 16px; color: #999; font-style: italic; }
-    </style>
-    """
-    for page in pages_json:
-        page_num = page.get("page", "?")
-        regions  = page.get("regions", [])
-        html += f'<div class="page-block">'
-        html += f'<div class="page-title">📄 Page {page_num}</div>'
-        if not regions:
-            html += '<div class="empty-page">Empty page</div>'
-        else:
-            for region in regions:
-                label   = region.get("label", "text")
-                content = region.get("content", "")
-                source  = region.get("source", "")
-                label_class = {"header": "label-header", "footer": "label-footer",
-                               "table": "label-table", "text": "label-text"}.get(label, "label-other")
-                source_text = "(PDF text layer)" if source == "pdf_text_layer" else "(API)"
-                html += f'<div class="region">'
-                html += f'<span class="region-label {label_class}">{label.upper()}</span>'
-                html += f'<span class="source-badge">{source_text}</span>'
-                html += f'<div class="region-content">{content}</div>'
-                html += f'</div>'
-        html += '</div>'
-    return html
-
-# ── STEP 6: Main OCR function ─────────────────────────────────
+            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 function ─────────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
-        return "Please upload a file.", "No regions detected.", "<p>No output yet.</p>"
+        return "Please upload a file."
 
     try:
         from glmocr import parse
 
-        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:
             import fitz
-            doc = fitz.open(path)
+            doc          = fitz.open(path)
             page_images  = []
             page_heights = []
-            page_widths  = []
             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"
                 pix.save(img_path)
                 page_images.append(img_path)
-                # Store original page dimensions for bbox math
-                page_heights.append(doc[i].rect.height)
-                page_widths.append(doc[i].rect.width)
+                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 = []
-            page_widths  = []
             results = parse(path)
             if not isinstance(results, list):
                 results = [results]
 
-        total_headers  = 0
-        total_footers  = 0
-        all_pages_json = []
-        all_summary    = []
+        all_pages_md = []
 
         for page_num, page_result in enumerate(results):
-            page_summary = [f"── PAGE {page_num + 1} of {len(results)} ──"]
-            page_regions = []
-            index        = 0
+            page_md, regions = get_page_data(page_result)
+            parts = []
 
-            # Get API regions first so we can analyze bbox gaps
-            data, dtype = get_page_regions(page_result)
-            api_regions = data if dtype == "json" else []
-
-            # For PDFs: check if API missed a top or bottom zone
             if is_pdf and page_num < len(page_heights):
-                ph = page_heights[page_num]
-                pw = page_widths[page_num]
-
-                # Scale bbox_2d: images were rendered at 1.5x, so bbox coords
-                # are in image pixels. Convert back to PDF page fractions.
-                # API bbox_2d uses the image dimensions (1.5x scaled).
-                # We need fractions of original page height.
-                top_gap, bottom_gap = get_missed_zones(
-                    api_regions,
-                    ph * 1.5,   # image height = page height * 1.5
-                    pw * 1.5,
-                )
-
-                # Top gap = missed header zone
+                top_gap, bottom_gap = get_missed_zones(regions, page_heights[page_num])
+
+                # HEADER at TOP
                 if top_gap is not None:
                     hdr = extract_zone_text(path, page_num, 0, top_gap)
                     if hdr:
-                        total_headers += 1
-                        page_regions.append({
-                            "index":   index,
-                            "label":   "header",
-                            "content": hdr,
-                            "source":  "pdf_text_layer"
-                        })
-                        page_summary.append("🔵 HEADER: " + hdr[:80])
-                        index += 1
-
-                # Bottom gap = missed footer zone
+                        parts.append(hdr)
+
+                # BODY in MIDDLE
+                if page_md:
+                    parts.append(page_md)
+
+                # FOOTER at BOTTOM
                 if bottom_gap is not None:
                     ftr = extract_zone_text(path, page_num, bottom_gap, 1.0)
                     if ftr:
-                        total_footers += 1
-                        page_regions.append({
-                            "index":   index,
-                            "label":   "footer",
-                            "content": ftr,
-                            "source":  "pdf_text_layer"
-                        })
-                        page_summary.append("🟢 FOOTER: " + ftr[:80])
-                        index += 1
-
-            # Now add API regions
-            if dtype == "json":
-                for region in api_regions:
-                    if not isinstance(region, dict):
-                        continue
-                    label   = region.get("label", "text")
-                    content = str(region.get("content", ""))
-                    bbox    = region.get("bbox_2d", [])
-
-                    if label in ABANDON:
-                        continue
-
-                    entry = {
-                        "index":   index,
-                        "label":   label,
-                        "content": content,
-                        "source":  "api"
-                    }
-                    if bbox:
-                        entry["bbox_2d"] = bbox
-
-                    if label == "header":
-                        total_headers += 1
-                        page_summary.append("🔵 HEADER (API): " + content[:80])
-                    elif label == "footer":
-                        total_footers += 1
-                        page_summary.append("🟢 FOOTER (API): " + content[:80])
-                    else:
-                        page_summary.append("[" + label + "]: " + content[:80])
-
-                    page_regions.append(entry)
-                    index += 1
-
-            elif dtype == "raw":
-                page_regions.append({
-                    "index":   index,
-                    "label":   "text",
-                    "content": data,
-                    "source":  "api_raw"
-                })
-                page_summary.append("[raw markdown — " + str(len(data)) + " chars]")
-
+                        parts.append(ftr)
             else:
-                page_summary.append("[empty page]")
+                if page_md:
+                    parts.append(page_md)
 
-            all_pages_json.append({
-                "page":    page_num + 1,
-                "regions": page_regions
-            })
-            all_summary.extend(page_summary)
-            all_summary.append("")
+            if parts:
+                all_pages_md.append("\n\n".join(parts))
 
-        final_json = json.dumps(all_pages_json, indent=2, ensure_ascii=False)
-
-        summary = (
-            "Total pages   : " + str(len(results)) + "\n"
-            + "Headers found : " + str(total_headers) + "\n"
-            + "Footers found : " + str(total_footers) + "\n"
-            + "─" * 40 + "\n"
-            + "\n".join(all_summary)
-        )
-
-        readable = json_to_readable_html(all_pages_json)
-        return final_json, summary, readable
+        final = "\n\n---\n\n".join(all_pages_md) if all_pages_md else "(No content)"
+        return final
 
     except Exception as e:
         import traceback
-        err = "Error: " + str(e) + "\n\n" + traceback.format_exc()
-        return err, "Failed.", "<pre>" + err + "</pre>"
+        return "Error: " + str(e) + "\n\n" + traceback.format_exc()
+
 
-# ── STEP 7: Gradio UI ─────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — MaaS") as demo:
+# ── STEP 5: Gradio UI ─────────────────────────────────────────
+with gr.Blocks(title="GLM-OCR") as demo:
     gr.Markdown("""
-    # 🔍 GLM-OCR — Zhipu MaaS
-    Upload PDF or image. Headers 🔵 and Footers 🟢 detected automatically.
-    Works for any PDF — no hard-coded assumptions.
-    Multi-page PDFs fully supported.
+    # 🔍 GLM-OCR
+    Upload a PDF or image to extract text.
+    Headers and footers included in correct positions.
     """)
 
     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.Tabs():
-        with gr.Tab("📋 Readable View"):
-            readable_out = gr.HTML(label="")
-        with gr.Tab("{ } JSON Output"):
-            json_out = gr.Textbox(lines=40, label="Raw JSON")
-        with gr.Tab("🗂️ Detected Regions"):
-            regions_out = gr.Textbox(lines=30, label="")
-
-    run_btn.click(
-        fn=run_ocr,
-        inputs=file_input,
-        outputs=[json_out, regions_out, readable_out]
-    )
+    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()