Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| import logging | |
| import os | |
| import re | |
| import tempfile | |
| import uuid | |
| from typing import List, Tuple | |
| import yaml | |
| try: | |
| import glmocr | |
| GLMOCR_BASE = os.path.dirname(glmocr.__file__) | |
| CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml") | |
| except ImportError: | |
| glmocr = None | |
| GLMOCR_BASE = "" | |
| CONFIG_PATH = "" | |
| log = logging.getLogger("glmocr_simple_app") | |
| logging.basicConfig(level=logging.INFO) | |
| GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito" | |
| if not GLMOCR_API_KEY: | |
| log.warning( | |
| "No ZHIPU_API_KEY or GLMOCR_API_KEY in environment; GlmOcr() will fail until you set one." | |
| ) | |
| RENDER_SCALE = 3.0 | |
| PAD_LEFT_FRAC = 0.035 | |
| PAD_RIGHT_FRAC = 0.10 | |
| PAD_TOP_FRAC = 0.018 | |
| PAD_BOTTOM_FRAC = 0.018 | |
| ENABLE_CONTRAST = True | |
| CONTRAST_FACTOR = 1.18 | |
| ENABLE_UNSHARP = True | |
| UNSHARP_RADIUS = 0.78 | |
| UNSHARP_PERCENT = 76 | |
| UNSHARP_THRESHOLD = 1 | |
| DEFAULT_ZONE_FRAC = 0.12 | |
| PDF_HEADER_BAND_FRAC = 0.10 | |
| ENABLE_FOOTER_OCR = True | |
| PDF_FOOTER_BAND_FRAC = 0.88 | |
| MIN_CROP_HEIGHT = 112 | |
| MIN_CROP_PIXELS = 112 * 112 | |
| PAGE_PNG_COMPRESS_LEVEL = 3 | |
| ZONE_JPEG_QUALITY = 95 | |
| _parser = None | |
| def _enhance_raster_for_ocr(img): | |
| from PIL import ImageEnhance, ImageFilter | |
| if ENABLE_CONTRAST: | |
| img = ImageEnhance.Contrast(img).enhance(CONTRAST_FACTOR) | |
| if ENABLE_UNSHARP: | |
| img = img.filter( | |
| ImageFilter.UnsharpMask( | |
| radius=UNSHARP_RADIUS, | |
| percent=UNSHARP_PERCENT, | |
| threshold=UNSHARP_THRESHOLD, | |
| ) | |
| ) | |
| return img | |
| def get_parser(): | |
| global _parser | |
| if glmocr is None: | |
| raise RuntimeError("glmocr is not installed.") | |
| if _parser is None: | |
| from glmocr import GlmOcr | |
| kw = {"mode": "maas"} | |
| if GLMOCR_API_KEY: | |
| kw["api_key"] = GLMOCR_API_KEY | |
| _parser = GlmOcr(**kw) | |
| return _parser | |
| if CONFIG_PATH: | |
| try: | |
| with open(CONFIG_PATH, "r", encoding="utf-8") as f: | |
| config = yaml.safe_load(f) | |
| config.setdefault("pipeline", {}).setdefault("maas", {}) | |
| config["pipeline"]["maas"]["enabled"] = True | |
| config["pipeline"]["maas"]["api_key"] = GLMOCR_API_KEY | |
| with open(CONFIG_PATH, "w", encoding="utf-8") as f: | |
| yaml.dump(config, f, default_flow_style=False, sort_keys=False) | |
| except Exception: | |
| pass | |
| 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 extract_pdf_text_in_band(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 = page.rect.height | |
| y_lo = h * y_start_frac | |
| y_hi = h * y_end_frac | |
| words = page.get_text("words") | |
| doc.close() | |
| parts = [] | |
| for w in words: | |
| if len(w) >= 5: | |
| y0, y1 = float(w[1]), float(w[3]) | |
| if y0 < y_hi and y1 > y_lo: | |
| parts.append(w[4]) | |
| return " ".join(parts).strip() | |
| except Exception: | |
| return "" | |
| def ocr_zone(image_path, y_start_frac, y_end_frac): | |
| 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: | |
| return "" | |
| crop = img.crop((0, y0, w, y1)) | |
| cw, ch = crop.size | |
| 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)) | |
| else: | |
| canvas.paste(crop, (0, need_h - ch)) | |
| crop = canvas | |
| fd, path = tempfile.mkstemp(suffix=".jpg") | |
| os.close(fd) | |
| try: | |
| crop.save(path, "JPEG", quality=ZONE_JPEG_QUALITY) | |
| parser = get_parser() | |
| out = parser.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 as e: | |
| log.warning("[%s] ocr_zone failed: %s", zone_name, e, exc_info=True) | |
| return "" | |
| def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]: | |
| import pymupdf as fitz | |
| from PIL import Image | |
| doc = fitz.open(pdf_path) | |
| page_images: List[str] = [] | |
| page_heights: List[int] = [] | |
| for i in range(len(doc)): | |
| page = doc[i] | |
| pix = page.get_pixmap(matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False) | |
| img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) | |
| img = _enhance_raster_for_ocr(img) | |
| w, h = img.size | |
| pad_l = int(w * PAD_LEFT_FRAC) | |
| pad_r = int(w * PAD_RIGHT_FRAC) | |
| pad_t = int(h * PAD_TOP_FRAC) | |
| pad_b = int(h * PAD_BOTTOM_FRAC) | |
| if any(p > 0 for p in (pad_l, pad_r, pad_t, pad_b)): | |
| canvas = Image.new("RGB", (w + pad_l + pad_r, h + pad_t + pad_b), (255, 255, 255)) | |
| canvas.paste(img, (pad_l, pad_t)) | |
| img = canvas | |
| uniq = uuid.uuid4().hex[:10] | |
| img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{uniq}_{i}.png") | |
| img.save(img_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL) | |
| page_images.append(img_path) | |
| page_heights.append(img.height) | |
| doc.close() | |
| return page_images, page_heights | |
| 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." | |
| page_images: List[str] = [] | |
| try: | |
| path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file) | |
| is_pdf = path.lower().endswith(".pdf") | |
| parser = get_parser() | |
| page_heights: List[int] = [] | |
| if is_pdf: | |
| page_images, page_heights = render_pdf_pages_to_images(path) | |
| results = parser.parse(page_images) | |
| else: | |
| page_images = [path] | |
| page_heights = [1000] | |
| 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) | |
| img_h = page_heights[page_num] if page_num < len(page_heights) else 1000 | |
| header_end_frac, footer_start_frac = get_header_footer_zones(regions, img_h) | |
| 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) | |
| he = max(0.02, min(0.25, he)) | |
| fs = max(0.75, min(0.98, fs)) | |
| parts = [] | |
| hdr = "" | |
| if is_pdf: | |
| hdr = extract_zone_text_pdf(path, page_num, 0, he) | |
| if not (hdr and hdr.strip()): | |
| hdr = extract_pdf_text_in_band(path, page_num, 0, PDF_HEADER_BAND_FRAC) | |
| if not (hdr and hdr.strip()) and page_num < len(page_images): | |
| hdr = ocr_zone(page_images[page_num], 0, he) | |
| if hdr and hdr.strip(): | |
| parts.append(hdr.strip()) | |
| if page_md and page_md.strip(): | |
| parts.append(page_md.strip()) | |
| if ENABLE_FOOTER_OCR and page_num < len(page_images): | |
| ftr = "" | |
| if is_pdf: | |
| ftr = extract_zone_text_pdf(path, page_num, fs, 1.0) | |
| if not (ftr and ftr.strip()): | |
| ftr = extract_pdf_text_in_band(path, page_num, PDF_FOOTER_BAND_FRAC, 1.0) | |
| if not (ftr and ftr.strip()): | |
| ftr = ocr_zone(page_images[page_num], fs, 1.0) | |
| if ftr and ftr.strip(): | |
| ftr_clean = ftr.strip() | |
| ftr_first_line = next( | |
| (ln.strip().lower() for ln in ftr_clean.splitlines() if ln.strip()), | |
| "", | |
| ) | |
| already_present = ftr_first_line and any( | |
| ftr_first_line in part.lower() for part in parts | |
| ) | |
| _footer_date_re = re.compile(r"\b\d{1,2}[-/]\d{2}\b") | |
| _footer_amt_re = re.compile(r"\b\d{1,3}(?:,\d{3})*\.\d{2}\b") | |
| _date_hits = len(_footer_date_re.findall(ftr_clean)) | |
| _amt_hits = len(_footer_amt_re.findall(ftr_clean)) | |
| is_txn_dump = _date_hits >= 3 and _amt_hits >= 3 | |
| if not already_present and not is_txn_dump: | |
| parts.append(ftr_clean) | |
| if parts: | |
| all_pages.append("\n\n".join(parts)) | |
| merged = "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)" | |
| return merged | |
| except Exception as e: | |
| import traceback | |
| log.exception("run_ocr failed: %s", e) | |
| return f"Error: {e}\n\n{traceback.format_exc()}" | |
| finally: | |
| for p in page_images: | |
| try: | |
| if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p): | |
| os.unlink(p) | |
| except Exception: | |
| pass | |
| def _create_gradio_demo(): | |
| import gradio as gr | |
| with gr.Blocks(title="GLM-OCR (simple)") as demo: | |
| gr.Markdown("# GLM-OCR") | |
| 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 (markdown / HTML from model)") | |
| run_btn.click(fn=run_ocr, inputs=file_in, outputs=out) | |
| return demo | |
| if __name__ == "__main__": | |
| _create_gradio_demo().launch() | |