app.py

"""
GLM-OCR Hugging Face Space app with client-side header/footer fallback for MaaS.
Works for every PDF and every image: uses API bboxes when available, else minimal band.
"""
# Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
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 sys
import tempfile
import yaml
import gradio as gr

import glmocr

# Log to stderr so Hugging Face / Docker show it in logs. Set GLMOCR_LOG=DEBUG for more detail.
_log_level = os.environ.get("GLMOCR_LOG", "INFO").upper()
logging.basicConfig(
    level=getattr(logging, _log_level, logging.INFO),
    format="[%(levelname)s] %(name)s: %(message)s",
    stream=sys.stderr,
)
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")

# When MaaS is enabled, the cloud API ignores local layout config and does not return
# header/footer regions. We always run client-side OCR on top/bottom bands as fallback.
DEFAULT_ZONE_FRAC = float(os.environ.get("GLMOCR_DEFAULT_ZONE_FRAC", "0.12"))
# MaaS rejects very small images (400). Only send crops that meet minimum dimensions.
MIN_CROP_HEIGHT = int(os.environ.get("GLMOCR_MIN_CROP_HEIGHT", "112"))
MIN_CROP_PIXELS = int(os.environ.get("GLMOCR_MIN_CROP_PIXELS", "12544"))  # 112*112

# Single shared parser to avoid "GLM-OCR initialized" per request and asyncio cleanup issues.
_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

# ---------------------------------------------------------------------------
# 1. Config: set MaaS and optionally include headers/footers for non-MaaS
# ---------------------------------------------------------------------------
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_as_text = {"header", "footer"}
    layout_section = config.get("pipeline", {}).get("layout", {})
    if "label_task_mapping" in layout_section:
        mapping = layout_section["label_task_mapping"]
        abandon = mapping.get("abandon")
        if isinstance(abandon, list):
            mapping["abandon"] = [x for x in abandon if x not in to_include_as_text]
        text_labels = mapping.get("text")
        if isinstance(text_labels, list):
            for label in to_include_as_text:
                if label not in text_labels:
                    text_labels.append(label)

    formatter_section = config.get("pipeline", {}).get("result_formatter", {})
    if "abandon" in formatter_section:
        abandon = formatter_section["abandon"]
        if isinstance(abandon, list):
            formatter_section["abandon"] = [x for x in abandon if x not in to_include_as_text]
    if "label_visualization_mapping" in formatter_section:
        text_vis = formatter_section["label_visualization_mapping"].get("text")
        if isinstance(text_vis, list):
            for label in to_include_as_text:
                if label not in text_vis:
                    text_vis.append(label)

    with open(CONFIG_PATH, "w") as f:
        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
except Exception:
    pass  # e.g. read-only package on Hugging Face; MaaS ignores layout anyway

# ---------------------------------------------------------------------------
# 2. result_formatter: stop stripping header/footer (best-effort; may be read-only)
# ---------------------------------------------------------------------------
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


def get_header_footer_zones(regions, norm_height=1000):
    if not regions:
        log.debug("get_header_footer_zones: no regions")
        return None, None
    y_tops, y_bottoms = [], []
    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])
            y_bottoms.append(bbox[3])
    if not y_tops:
        log.debug("get_header_footer_zones: no bboxes in regions")
        return None, None
    he_frac = min(y_tops) / norm_height
    fs_frac = max(y_bottoms) / norm_height
    log.debug("get_header_footer_zones: regions=%s -> header_end_frac=%.3f footer_start_frac=%.3f", len(regions), he_frac, fs_frac)
    return he_frac, fs_frac


def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
    try:
        import pymupdf as fitz
        doc = fitz.open(pdf_path)
        page = doc[page_num]
        h, w = page.rect.height, 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()
        log.debug("extract_zone_text_pdf page=%s y=%.2f-%.2f -> %d chars", page_num, y_start_frac, y_end_frac, len(text))
        return text
    except Exception as e:
        log.debug("extract_zone_text_pdf failed: %s", e)
        return ""


def ocr_zone(image_path, y_start_frac, y_end_frac):
    """Run OCR on a horizontal band. Pads small crops to meet API minimum size instead of skipping."""
    zone_name = "header" if y_end_frac < 0.5 else "footer"
    try:
        from PIL import Image
        img = Image.open(image_path).convert("RGB")
        w, h = img.size
        y0 = max(0, int(h * y_start_frac))
        y1 = min(h, int(h * y_end_frac))
        if y1 <= y0:
            log.info("[%s] ocr_zone: skip (zero height band y0=%s y1=%s)", zone_name, y0, y1)
            return ""
        crop = img.crop((0, y0, w, y1))
        cw, ch = crop.size
        # Pad to minimum size so MaaS accepts the image (avoids 400 on tiny strips)
        if ch < MIN_CROP_HEIGHT or (cw * ch) < MIN_CROP_PIXELS:
            need_h = max(ch, MIN_CROP_HEIGHT)
            need_w = max(cw, 1)
            if (need_w * need_h) < MIN_CROP_PIXELS:
                need_w = max(need_w, (MIN_CROP_PIXELS + need_h - 1) // need_h)
            canvas = Image.new("RGB", (need_w, need_h), (255, 255, 255))
            if zone_name == "header":
                canvas.paste(crop, (0, 0))  # crop at top
            else:
                canvas.paste(crop, (0, need_h - ch))  # crop at bottom
            crop = canvas
            cw, ch = crop.size
            log.info("[%s] ocr_zone: padded to %dx%d for API", zone_name, cw, ch)
        fd, path = tempfile.mkstemp(suffix=".jpg")
        os.close(fd)
        try:
            crop.save(path, "JPEG", quality=92)
            parser = get_parser()
            out = parser.parse(path)
            if not isinstance(out, list):
                out = [out]
            if out and getattr(out[0], "markdown_result", None):
                text = (out[0].markdown_result or "").strip()
                log.info("[%s] ocr_zone: crop %dx%d -> %d chars", zone_name, cw, ch, len(text))
                return text
            log.info("[%s] ocr_zone: crop %dx%d -> empty result from API", zone_name, cw, ch)
        finally:
            try:
                os.unlink(path)
            except Exception:
                pass
    except Exception as e:
        log.warning("[%s] ocr_zone failed: %s", zone_name, e, exc_info=True)
    return ""


def get_page_md_and_regions(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, dict) and "regions" in jr:
            regions = jr.get("regions") or []
        elif 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(r, dict) and "regions" in r:
                regions = r.get("regions") or []
    return md, regions


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")
        log.info("run_ocr: path=%s is_pdf=%s", path, is_pdf)
        parser = get_parser()

        if is_pdf:
            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 = os.path.join(tempfile.gettempdir(), f"maas_page_{i}.png")
                pix.save(img_path)
                page_images.append(img_path)
            doc.close()
            log.info("run_ocr: PDF has %d pages, running MaaS on page images", len(page_images))
            results = parser.parse(page_images)
        else:
            page_images = [path]
            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_md_and_regions(page_result)
            header_end_frac, footer_start_frac = get_header_footer_zones(regions, 1000)

            # MaaS does not return header/footer regions; always use at least default bands
            he = header_end_frac if header_end_frac is not None else DEFAULT_ZONE_FRAC
            fs = footer_start_frac if footer_start_frac is not None else (1.0 - DEFAULT_ZONE_FRAC)
            if he <= 0:
                he = DEFAULT_ZONE_FRAC  # ensure we always OCR a top band
            if fs >= 1.0:
                fs = 1.0 - DEFAULT_ZONE_FRAC  # ensure we always OCR a bottom band

            log.info(
                "run_ocr: page %d regions=%s he=%.3f fs=%.3f body_md_len=%d",
                page_num, len(regions), he, fs, len(page_md or ""),
            )

            parts = []

            # Always run header band (top 8–12% of page)
            if page_num < len(page_images):
                img_path = page_images[page_num]
                hdr = ""
                if is_pdf:
                    hdr = extract_zone_text_pdf(path, page_num, 0, he)
                    if hdr and hdr.strip():
                        log.info("run_ocr: page %d header from PDF extract: %d chars", page_num, len(hdr.strip()))
                if not (hdr and hdr.strip()):
                    hdr = ocr_zone(img_path, 0, he)
                if hdr and hdr.strip():
                    parts.append(hdr.strip())
                elif not (hdr and hdr.strip()) and he > 0:
                    log.info("run_ocr: page %d header empty (PDF extract + ocr_zone)", page_num)

            if page_md:
                parts.append(page_md)

            # Always run footer band (bottom 8–12% of page)
            if page_num < len(page_images):
                img_path = page_images[page_num]
                ftr = ""
                if is_pdf:
                    ftr = extract_zone_text_pdf(path, page_num, fs, 1.0)
                    if ftr and ftr.strip():
                        log.info("run_ocr: page %d footer from PDF extract: %d chars", page_num, len(ftr.strip()))
                if not (ftr and ftr.strip()):
                    ftr = ocr_zone(img_path, fs, 1.0)
                if ftr and ftr.strip():
                    parts.append(ftr.strip())
                elif not (ftr and ftr.strip()) and fs < 1.0:
                    log.info("run_ocr: page %d footer empty (PDF extract + ocr_zone)", page_num)

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

        log.info("run_ocr: done, %d pages in output", len(all_pages))
        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()}"


with gr.Blocks(title="GLM-OCR") as demo:
    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers and footers included.")
    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()