app.py

"""
GLM-OCR Hugging Face Space — MaaS mode.
Extracts header from top zone the API missed. No footer extraction.
"""
import asyncio
try:
    _orig_close = asyncio.BaseEventLoop.close
    def _safe_close(self):
        try:
            _orig_close(self)
        except (ValueError, OSError):
            pass
    asyncio.BaseEventLoop.close = _safe_close
except Exception:
    pass

import logging
import os
import re
import tempfile
import yaml
import gradio as gr
import glmocr

log = logging.getLogger("glmocr_app")

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

# ── Fix config ────────────────────────────────────────────────
try:
    with open(CONFIG_PATH, "r") as f:
        config = yaml.safe_load(f)
    config["pipeline"]["maas"]["enabled"] = True
    config["pipeline"]["maas"]["api_key"] = os.environ.get(
        "GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
    )
    to_include = {"header", "footer"}
    layout_section = config.get("pipeline", {}).get("layout", {})
    if "label_task_mapping" in layout_section:
        mapping = layout_section["label_task_mapping"]
        if isinstance(mapping.get("abandon"), list):
            mapping["abandon"] = [x for x in mapping["abandon"] if x not in to_include]
        if isinstance(mapping.get("text"), list):
            for lb in to_include:
                if lb not in mapping["text"]:
                    mapping["text"].append(lb)
    formatter_section = config.get("pipeline", {}).get("result_formatter", {})
    if isinstance(formatter_section.get("abandon"), list):
        formatter_section["abandon"] = [x for x in formatter_section["abandon"] if x not in to_include]
    if "label_visualization_mapping" in formatter_section:
        tv = formatter_section["label_visualization_mapping"].get("text")
        if isinstance(tv, list):
            for lb in to_include:
                if lb not in tv:
                    tv.append(lb)
    with open(CONFIG_PATH, "w") as f:
        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
except Exception:
    pass

# ── Fix result_formatter ──────────────────────────────────────
try:
    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'"):
        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)
except Exception:
    pass

# ── Shared parser ─────────────────────────────────────────────
_parser = None
def get_parser():
    global _parser
    if _parser is None:
        from glmocr import GlmOcr
        _parser = GlmOcr(
            api_key=os.environ.get("GLMOCR_API_KEY", "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"),
            mode="maas",
        )
    return _parser

# ── Helpers ───────────────────────────────────────────────────
def get_top_gap_frac(regions, img_height):
    """
    Returns the fraction of page height that the API missed at the top.
    Based on the y coordinate of the topmost bbox_2d region.
    Returns None if no gap detected (API covered the full page top).
    """
    y_tops = []
    for r in regions:
        bbox = r.get("bbox_2d") if isinstance(r, dict) else getattr(r, "bbox_2d", None)
        if bbox and len(bbox) >= 4:
            y_tops.append(bbox[1])
    if not y_tops:
        # No bbox info from API — default to extracting top 12%
        return 0.12
    first_y = min(y_tops)
    # Only treat as a missed header if gap is more than 8% of page
    return (first_y / img_height) if first_y > img_height * 0.08 else None


def extract_header_pdf(pdf_path, page_num, y_end_frac):
    """Extract text from top zone of PDF page using clip rect."""
    try:
        import pymupdf as fitz
        doc  = fitz.open(pdf_path)
        page = doc[page_num]
        h, w = page.rect.height, page.rect.width
        text = page.get_text(clip=fitz.Rect(0, 0, w, h * y_end_frac)).strip()
        doc.close()
        return text
    except Exception:
        return ""


def get_page_data(page_result):
    md = ""
    if hasattr(page_result, "markdown_result") and page_result.markdown_result:
        md = (page_result.markdown_result or "").strip()
    regions = []
    if hasattr(page_result, "json_result"):
        jr = page_result.json_result
        if isinstance(jr, list) and len(jr) > 0:
            r = jr[0] if isinstance(jr[0], list) else jr
            if isinstance(r, list):
                regions = r
        elif isinstance(jr, dict) and "regions" in jr:
            regions = jr.get("regions") or []
    return md, regions


# ── Main OCR ──────────────────────────────────────────────────
def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a file."
    try:
        import pymupdf as fitz

        path   = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
        is_pdf = path.lower().endswith(".pdf")
        parser = get_parser()

        if is_pdf:
            doc          = fitz.open(path)
            page_images  = []
            page_heights = []
            for i in range(len(doc)):
                pix      = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
                img_path = os.path.join(tempfile.gettempdir(), f"maas_page_{i}.png")
                pix.save(img_path)
                page_images.append(img_path)
                page_heights.append(doc[i].rect.height * 1.5)
            doc.close()
            results = parser.parse(page_images)
        else:
            page_images  = [path]
            page_heights = []
            results      = parser.parse(path)

        if not isinstance(results, list):
            results = [results]

        all_pages = []
        for page_num, page_result in enumerate(results):
            page_md, regions = get_page_data(page_result)
            parts = []

            # HEADER — extract top zone missed by API (PDF only)
            if is_pdf and page_num < len(page_heights):
                top_gap = get_top_gap_frac(regions, page_heights[page_num])
                if top_gap is not None:
                    hdr = extract_header_pdf(path, page_num, top_gap)
                    if hdr:
                        parts.append(hdr)

            # BODY — full API output (already includes real footer text)
            if page_md:
                parts.append(page_md)

            if parts:
                all_pages.append("\n\n".join(parts))

        return "\n\n---\n\n".join(all_pages) if all_pages else "(No content)"

    except Exception as e:
        import traceback
        log.exception("run_ocr failed: %s", e)
        return f"Error: {e}\n\n{traceback.format_exc()}"


# ── UI ────────────────────────────────────────────────────────
with gr.Blocks(title="GLM-OCR") as demo:
    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers included in correct position.")
    file_in = gr.File(label="Upload PDF or image",
                      file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
    run_btn = gr.Button("Run OCR", variant="primary")
    out     = gr.Textbox(lines=40, label="Output")
    run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)

if __name__ == "__main__":
    demo.launch()