Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import yaml
 import re
 import os
+import json
 
 import glmocr
 GLMOCR_BASE    = os.path.dirname(glmocr.__file__)
@@ -49,8 +50,6 @@ 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.
-# Headers and footers are extracted directly from the PDF using PyMuPDF.
 def extract_pdf_headers(path):
     import fitz
     page_headers = []
@@ -67,21 +66,13 @@ def extract_pdf_headers(path):
     return page_headers
 
 # ── STEP 4: Parse one page result into regions list ───────────
-# result from parse() is a LIST of PipelineResult objects (one per page)
-# Each PipelineResult has:
-#   .json_result  → list of region dicts with 'label' and 'content'
-#   .markdown_result → plain markdown string
 def get_page_regions(page_result):
-    # Try json_result first (structured regions)
     if hasattr(page_result, "json_result"):
         jr = page_result.json_result
         if isinstance(jr, list) and len(jr) > 0:
-            # json_result is a list of lists (one list per page image)
-            # since we pass one image per call, take index 0
             regions = jr[0] if isinstance(jr[0], list) else jr
             if isinstance(regions, list):
                 return regions, "json"
-    # Fallback to markdown_result
     if hasattr(page_result, "markdown_result"):
         md = page_result.markdown_result
         if md and isinstance(md, str) and md.strip():
@@ -109,7 +100,6 @@ def run_ocr(uploaded_file):
                 pix.save(img_path)
                 page_images.append(img_path)
             doc.close()
-            # parse() returns a LIST — one PipelineResult per page
             results = parse(page_images)
             if not isinstance(results, list):
                 results = [results]
@@ -121,27 +111,35 @@ def run_ocr(uploaded_file):
 
         total_headers = 0
         total_footers = 0
-        all_markdown  = []
-        all_summary   = []
+        all_pages_json = []   # final JSON output — list of pages
+        all_summary    = []
 
         for page_num, page_result in enumerate(results):
-            page_md      = []
-            page_summary = [f"── PAGE {page_num + 1} of {len(results)} ──"]
+            page_summary  = [f"── PAGE {page_num + 1} of {len(results)} ──"]
+            page_regions  = []   # JSON regions for this page
 
-            # Inject header/footer from PDF text layer
+            # Inject header from PDF text layer
             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_regions.append({
+                        "label":   "header",
+                        "content": hdr,
+                        "source":  "pdf_text_layer"
+                    })
                     page_summary.append("🔵 HEADER: " + hdr[:100])
                 if ftr:
                     total_footers += 1
-                    page_md.append("<!-- FOOTER -->\n" + ftr)
+                    page_regions.append({
+                        "label":   "footer",
+                        "content": ftr,
+                        "source":  "pdf_text_layer"
+                    })
                     page_summary.append("🟢 FOOTER: " + ftr[:100])
 
-            # Get regions from this page result
+            # Get regions from API result
             data, dtype = get_page_regions(page_result)
 
             if dtype == "json":
@@ -157,27 +155,49 @@ def run_ocr(uploaded_file):
                     if label == "header":
                         total_headers += 1
                         page_summary.append("🔵 HEADER (API): " + content[:100])
-                        page_md.append("<!-- HEADER -->\n" + content)
+                        page_regions.append({
+                            "label":   "header",
+                            "content": content,
+                            "source":  "api"
+                        })
                     elif label == "footer":
                         total_footers += 1
                         page_summary.append("🟢 FOOTER (API): " + content[:100])
-                        page_md.append("<!-- FOOTER -->\n" + content)
+                        page_regions.append({
+                            "label":   "footer",
+                            "content": content,
+                            "source":  "api"
+                        })
                     else:
                         page_summary.append("[" + label + "]: " + content[:100])
-                        page_md.append(content)
+                        page_regions.append({
+                            "label":   label,
+                            "content": content,
+                            "source":  "api"
+                        })
 
             elif dtype == "raw":
-                page_md.append(data)
-                page_summary.append("[markdown output]")
-                page_summary.append("Content length: " + str(len(data)) + " chars")
+                # Wrap raw markdown as a single text region
+                page_regions.append({
+                    "label":   "text",
+                    "content": data,
+                    "source":  "api_raw"
+                })
+                page_summary.append("[raw markdown — " + str(len(data)) + " chars]")
 
             else:
                 page_summary.append("[empty page]")
 
-            all_markdown.append("\n\n".join(page_md))
+            all_pages_json.append({
+                "page":    page_num + 1,
+                "regions": page_regions
+            })
             all_summary.extend(page_summary)
             all_summary.append("")
 
+        # Format final JSON output
+        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"
@@ -186,18 +206,18 @@ def run_ocr(uploaded_file):
             + "\n".join(all_summary)
         )
 
-        markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
-        return markdown, summary
+        return final_json, summary
 
     except Exception as e:
         import traceback
         return "Error: " + str(e) + "\n\n" + traceback.format_exc(), "Failed."
 
 # ── STEP 6: Gradio UI ─────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
+with gr.Blocks(title="GLM-OCR — MaaS (JSON Output)") as demo:
     gr.Markdown("""
     # 🔍 GLM-OCR — Zhipu MaaS
-    Upload PDF or image. Headers 🔵 and Footers 🟢 now appear in output.
+    Upload PDF or image. Output is **JSON format** with all regions.
+    Headers 🔵 and Footers 🟢 included in every page.
     Multi-page PDFs fully supported.
     """)
 
@@ -209,7 +229,7 @@ with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 📄 Markdown Output")
+            gr.Markdown("### 📄 JSON Output")
             markdown_out = gr.Textbox(lines=30, label="")
         with gr.Column():
             gr.Markdown("### 🗂️ Detected Regions")