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("\n" + hdr) page_summary.append("🔵 HEADER: " + hdr[:100]) if ftr: total_footers += 1 page_md.append("\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("\n" + content) elif label == "footer": total_footers += 1 page_summary.append("🟢 FOOTER (API): " + content[:100]) page_md.append("\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()