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