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