app.py

import gradio as gr
import yaml
import re
import os

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: 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",
]
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()

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)

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 = []
    doc = fitz.open(path)
    for page in doc:
        page_rect   = page.rect
        page_height = page_rect.height
        header_rect = fitz.Rect(0, 0, page_rect.width, page_height * 0.12)
        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: 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():
            return md.strip(), "raw"
    return [], "empty"

# ── STEP 5: Main OCR function ─────────────────────────────────
def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file.", "No regions detected."

    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
            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)
                img_path = f"/tmp/maas_page_{i}.png"
                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]
        else:
            pdf_headers = []
            results = parse(path)
            if not isinstance(results, list):
                results = [results]

        total_headers = 0
        total_footers = 0
        all_markdown  = []
        all_summary   = []

        for page_num, page_result in enumerate(results):
            page_md      = []
            page_summary = [f"── PAGE {page_num + 1} of {len(results)} ──"]

            # Inject header/footer 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_summary.append("🔵 HEADER: " + hdr[:100])
                if ftr:
                    total_footers += 1
                    page_md.append("<!-- FOOTER -->\n" + ftr)
                    page_summary.append("🟢 FOOTER: " + ftr[:100])

            # Get regions from this page result
            data, dtype = get_page_regions(page_result)

            if dtype == "json":
                for region in data:
                    if not isinstance(region, dict):
                        continue
                    label   = region.get("label", "text")
                    content = str(region.get("content", ""))

                    if label in ABANDON:
                        continue

                    if label == "header":
                        total_headers += 1
                        page_summary.append("🔵 HEADER (API): " + content[:100])
                        page_md.append("<!-- HEADER -->\n" + content)
                    elif label == "footer":
                        total_footers += 1
                        page_summary.append("🟢 FOOTER (API): " + content[:100])
                        page_md.append("<!-- FOOTER -->\n" + content)
                    else:
                        page_summary.append("[" + label + "]: " + content[:100])
                        page_md.append(content)

            elif dtype == "raw":
                page_md.append(data)
                page_summary.append("[markdown output]")
                page_summary.append("Content length: " + str(len(data)) + " chars")

            else:
                page_summary.append("[empty page]")

            all_markdown.append("\n\n".join(page_md))
            all_summary.extend(page_summary)
            all_summary.append("")

        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)
        )

        markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
        return markdown, 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:
    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.Row():
        with gr.Column():
            gr.Markdown("### 📄 Markdown Output")
            markdown_out = gr.Textbox(lines=30, label="")
        with gr.Column():
            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()