""" 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. """ import os import re import tempfile import yaml import gradio as gr 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") DEFAULT_ZONE_FRAC = float(os.environ.get("GLMOCR_DEFAULT_ZONE_FRAC", "0.08")) # --------------------------------------------------------------------------- # 1. Config: include headers/footers (remove from abandon lists) # --------------------------------------------------------------------------- with open(CONFIG_PATH, "r") as f: config = yaml.safe_load(f) config["pipeline"]["maas"]["enabled"] = True config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV" to_remove = {"header", "footer"} for key in ("layout", "result_formatter"): if key not in config.get("pipeline", {}): continue section = config["pipeline"][key] if key == "layout" and "label_task_mapping" in section: abandon = section["label_task_mapping"].get("abandon") if isinstance(abandon, list): section["label_task_mapping"]["abandon"] = [x for x in abandon if x not in to_remove] elif key == "result_formatter" and "abandon" in section: abandon = section["abandon"] if isinstance(abandon, list): section["abandon"] = [x for x in abandon if x not in to_remove] with open(CONFIG_PATH, "w") as f: yaml.dump(config, f, default_flow_style=False, sort_keys=False) # --------------------------------------------------------------------------- # 2. result_formatter: stop stripping header/footer in postprocess # --------------------------------------------------------------------------- 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) def get_header_footer_zones(regions, norm_height=1000): if not 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: return None, None return min(y_tops) / norm_height, max(y_bottoms) / norm_height 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() return text except Exception: return "" def ocr_zone(image_path, y_start_frac, y_end_frac): try: from PIL import Image from glmocr import parse 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: return "" crop = img.crop((0, y0, w, y1)) fd, path = tempfile.mkstemp(suffix=".png") os.close(fd) crop.save(path) try: out = parse(path) if not isinstance(out, list): out = [out] if out and getattr(out[0], "markdown_result", None): return (out[0].markdown_result or "").strip() finally: try: os.unlink(path) except Exception: pass except Exception: pass 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, 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: from glmocr import parse import pymupdf as fitz path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file) is_pdf = path.lower().endswith(".pdf") 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() results = parse(page_images) else: page_images = [path] results = 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) 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 and header_end_frac is not None: he = 0 if fs >= 1.0 and footer_start_frac is not None: fs = 1.0 parts = [] if he > 0 and 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 not (hdr and hdr.strip()): hdr = ocr_zone(img_path, 0, he) if hdr and hdr.strip(): parts.append(hdr.strip()) if page_md: parts.append(page_md) if fs < 1.0 and 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 not (ftr and ftr.strip()): ftr = ocr_zone(img_path, fs, 1.0) if ftr and ftr.strip(): parts.append(ftr.strip()) 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 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()