Spaces:
Running on Zero
Running on Zero
| """MiniCPM-OCR document parsing for a Hugging Face ZeroGPU Space. | |
| Both models are PyTorch. ZeroGPU only lends a GPU inside the Gradio function | |
| decorated with @spaces.GPU, so layout and recognition happen in that one call. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| # ZeroGPU patches torch at import time. Import it before torch on a Space, | |
| # and keep a no-op decorator for a normal local process. | |
| if os.environ.get("SPACE_ID"): | |
| import spaces | |
| else: | |
| class spaces: # noqa: N801 | |
| def GPU(*_args, **_kwargs): | |
| def decorate(fn): | |
| return fn | |
| return decorate | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from transformers import ( | |
| AutoImageProcessor, | |
| AutoModelForImageClassification, | |
| AutoModelForObjectDetection, | |
| AutoProcessor, | |
| MiniCPMOcrForConditionalGeneration, | |
| ) | |
| from layout import detect as detect_layout | |
| from postprocess import format_block | |
| OCR_ID = os.environ.get("OCR_MODEL_ID", "openbmb/MiniCPM-OCR-1.0") | |
| LAYOUT_ID = os.environ.get("LAYOUT_MODEL_ID", "PaddlePaddle/PP-DocLayoutV3_safetensors") | |
| ORI_ID = os.environ.get("ORI_MODEL_ID", "PaddlePaddle/PP-LCNet_x1_0_doc_ori_safetensors") | |
| PROMPTS = { | |
| "text": ("Text Recognition:", 2048), | |
| "content": ("Text Recognition:", 2048), | |
| "abstract": ("Text Recognition:", 2048), | |
| "reference": ("Text Recognition:", 1024), | |
| "reference_content": ("Text Recognition:", 1024), | |
| "vertical_text": ("Text Recognition:", 1024), | |
| "vision_footnote": ("Text Recognition:", 512), | |
| "algorithm": ("Text Recognition:", 2048), | |
| "doc_title": ("Text Recognition:", 256), | |
| "paragraph_title": ("Text Recognition:", 256), | |
| "figure_title": ("Text Recognition:", 256), | |
| "table": ("Table Recognition:", 2048), | |
| "formula": ("Formula Recognition:", 1024), | |
| "display_formula": ("Formula Recognition:", 1024), | |
| "inline_formula": ("Text Recognition:", 512), | |
| "formula_number": ("Text Recognition:", 128), | |
| "header": ("Text Recognition:", 256), | |
| "footer": ("Text Recognition:", 256), | |
| "footnote": ("Text Recognition:", 512), | |
| "aside_text": ("Text Recognition:", 512), | |
| "number": ("Text Recognition:", 128), | |
| "chart": ("Chart to Table:", 2048), | |
| "seal": ("Seal Recognition:", 256), | |
| } | |
| SKIP_MARKDOWN = {"number", "footnote", "header", "footer", "aside_text", "formula_number"} | |
| SKIP_VLM = {"image", "header_image", "footer_image"} | |
| COPY = { | |
| "en": { | |
| "language_label": "Language", | |
| "intro": "# Document parsing\nUpload an image.", | |
| "limit_note": ( | |
| 'Because of ZeroGPU limits, this Space can only run Transformers serially on a single image. ' | |
| 'For a smoother, more complete experience, use the <a href="https://platform.modelbest.ai/console/playground/document-parse" target="_blank" rel="noopener noreferrer">official demo</a>.' | |
| ), | |
| "image": "Image", | |
| "upload": "Drop an image here\n- or -\nclick to upload", | |
| "upright": "Auto-rotate", | |
| "upright_info": ( | |
| "Corrects a small skew and turns the page to 0°, 90°, 180°, or 270° " | |
| "before layout detection and recognition." | |
| ), | |
| "seal": "Recognize seals", | |
| "seal_info": "Reads the text inside seals. Turn off to skip them.", | |
| "chart": "Recognize charts", | |
| "chart_info": "Converts charts into tables. Turn off to skip them.", | |
| "parse": "Parse", | |
| "result": "Recognition result", | |
| "view_label": "Display", | |
| "rendered": "Rendered", | |
| "source": "Markdown", | |
| "parsing": "Parsing…", | |
| "need_image": "Please upload an image.", | |
| "empty": "No content was recognized on this page.", | |
| }, | |
| "zh": { | |
| "language_label": "语言", | |
| "intro": "# 文档解析\n上传一张图片。", | |
| "limit_note": ( | |
| '由于 ZeroGPU 的限制，这里只能用 Transformers 串行、单图推理。需要更流畅、更全面的体验，' | |
| '请前往<a href="https://platform.modelbest.ai/console/playground/document-parse" target="_blank" rel="noopener noreferrer">官方 Demo</a>。' | |
| ), | |
| "image": "图片", | |
| "upload": "将图片拖到这里\n- 或 -\n点击上传", | |
| "upright": "自适应旋转", | |
| "upright_info": "先做小角度摆正和 0°/90°/180°/270° 转向，再做版面检测和识别。", | |
| "seal": "印章识别", | |
| "seal_info": "识别印章里的文字。关闭后跳过印章。", | |
| "chart": "图表识别", | |
| "chart_info": "把图表转成表格。关闭后跳过图表。", | |
| "parse": "解析", | |
| "result": "识别结果", | |
| "view_label": "展示", | |
| "rendered": "渲染", | |
| "source": "原始 Markdown", | |
| "parsing": "正在解析…", | |
| "need_image": "请上传一张图片。", | |
| "empty": "这一页没有识别出内容。", | |
| }, | |
| } | |
| RESULT_CSS = """ | |
| .ocr-columns { align-items: flex-start !important; gap: 16px; } | |
| .ocr-columns > .column { flex: 1 1 0 !important; min-width: 0 !important; width: calc(50% - 8px) !important; } | |
| .ocr-col > .block, .ocr-col > .form { width: 100% !important; } | |
| .html:has(.ocr-note), | |
| .html-container:has(.ocr-note) { padding: 0 !important; } | |
| .ocr-note { margin: 0 0 12px; padding: 0; color: #6b7280; font-size: 13px; line-height: 1.5; } | |
| .ocr-note a { color: #2563eb; margin: 0; padding: 0; } | |
| .ocr-result table { border-collapse: collapse; width: 100%; margin: 8px 0; } | |
| .ocr-result td, .ocr-result th { border: 1px solid #d1d5db; padding: 4px 8px; vertical-align: top; } | |
| """ | |
| layout_model = AutoModelForObjectDetection.from_pretrained(LAYOUT_ID, dtype=torch.float32).to("cuda").eval() | |
| ori_processor = AutoImageProcessor.from_pretrained(ORI_ID) | |
| ori_model = AutoModelForImageClassification.from_pretrained(ORI_ID, dtype=torch.float32).to("cuda").eval() | |
| ocr_processor = AutoProcessor.from_pretrained(OCR_ID) | |
| ocr_model = MiniCPMOcrForConditionalGeneration.from_pretrained(OCR_ID, dtype=torch.bfloat16).to("cuda").eval() | |
| def _copy(language: str) -> dict: | |
| return COPY["zh"] if language == "zh" else COPY["en"] | |
| def _note(language: str) -> str: | |
| return f'<p class="ocr-note">{_copy(language)["limit_note"]}</p>' | |
| def _duration(image, upright, recognize_seal, recognize_chart, language): | |
| del image, upright, recognize_seal, recognize_chart, language | |
| return 90 | |
| def _rotate(image: Image.Image, angle: float) -> Image.Image: | |
| if abs(angle) < 0.3 or abs(abs(angle) - 360) < 0.3: | |
| return image | |
| nearest = round(angle / 90.0) * 90.0 | |
| if abs(angle - nearest) < 0.3: | |
| turns = int(nearest) % 360 | |
| return image.rotate(turns, expand=True, fillcolor=(255, 255, 255)) | |
| return image.rotate(angle, expand=True, fillcolor=(255, 255, 255), resample=Image.Resampling.BICUBIC) | |
| def _projection_score(binary: np.ndarray) -> float: | |
| profile = binary.sum(axis=1).astype(np.float64) | |
| mean = float(profile.mean()) | |
| if mean < 1e-6: | |
| return 0.0 | |
| delta = np.diff(profile) | |
| return float((delta * delta).mean() / (mean * mean)) | |
| def _deskew_angle(image: Image.Image) -> float: | |
| import cv2 | |
| gray = np.asarray(image.convert("L")) | |
| long_side = max(gray.shape) | |
| if long_side > 900: | |
| scale = 900 / long_side | |
| gray = cv2.resize( | |
| gray, | |
| (max(1, int(gray.shape[1] * scale)), max(1, int(gray.shape[0] * scale))), | |
| interpolation=cv2.INTER_AREA, | |
| ) | |
| block = max(15, (min(gray.shape) // 35) | 1) | |
| binary = cv2.adaptiveThreshold( | |
| gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, block, 15 | |
| ) | |
| def score(angle: float) -> float: | |
| height, width = binary.shape | |
| matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0) | |
| rotated = cv2.warpAffine(binary, matrix, (width, height), flags=cv2.INTER_NEAREST, borderValue=0) | |
| return _projection_score(rotated) | |
| angles = np.linspace(-15, 15, 31) | |
| scores = [score(float(angle)) for angle in angles] | |
| best = float(angles[int(np.argmax(scores))]) | |
| straight = scores[len(scores) // 2] | |
| if abs(best) <= 0.5 and straight >= max(scores) * 0.98: | |
| return 0.0 | |
| fine = np.linspace(best - 1.0, best + 1.0, 11) | |
| fine_scores = [score(float(angle)) for angle in fine] | |
| return float(fine[int(np.argmax(fine_scores))]) | |
| def _orientation_degrees(image: Image.Image) -> int: | |
| preview = image | |
| long_side = max(image.size) | |
| if long_side > 1600: | |
| scale = 1600 / long_side | |
| preview = image.resize((max(1, int(image.width * scale)), max(1, int(image.height * scale)))) | |
| inputs = ori_processor(images=preview, return_tensors="pt") | |
| inputs = {key: value.to("cuda") for key, value in inputs.items()} | |
| with torch.inference_mode(): | |
| label_id = int(ori_model(**inputs).last_hidden_state.argmax(-1).item()) | |
| label = str(ori_model.config.id2label[label_id]) | |
| digits = "".join(char for char in label if char.isdigit()) | |
| return int(digits or "0") % 360 | |
| def _upright(image: Image.Image) -> Image.Image: | |
| image = image.convert("RGB") | |
| skew = _deskew_angle(image) | |
| leveled = _rotate(image, skew) | |
| quadrant = _orientation_degrees(leveled) | |
| return _rotate(leveled, float(quadrant)) | |
| def _crop(image: Image.Image, box): | |
| width, height = image.size | |
| x1 = max(0, int(box[0])) | |
| y1 = max(0, int(box[1])) | |
| x2 = min(width, int(box[2]) + 1) | |
| y2 = min(height, int(box[3]) + 1) | |
| if x2 - x1 < 2 or y2 - y1 < 2: | |
| return None | |
| return image.crop((x1, y1, x2, y2)) | |
| def _strip_fence(text: str) -> str: | |
| lines = text.strip().splitlines() | |
| if len(lines) >= 2 and lines[0].startswith("```") and lines[-1].startswith("```"): | |
| return "\n".join(lines[1:-1]).strip() | |
| return text.strip() | |
| def _recognize(image: Image.Image, prompt: str, max_new_tokens: int) -> str: | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image}, | |
| {"type": "text", "text": prompt}, | |
| ], | |
| } | |
| ] | |
| inputs = ocr_processor.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| enable_thinking=False, | |
| processor_kwargs={"max_slice_nums": 9, "downsample_mode": "4x"}, | |
| ).to("cuda", dtype=torch.bfloat16) | |
| with torch.inference_mode(): | |
| output = ocr_model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=False, | |
| repetition_penalty=1.0, | |
| no_repeat_ngram_size=100, | |
| ) | |
| prompt_len = int(inputs["input_ids"].shape[1]) | |
| return _strip_fence(ocr_processor.decode(output[0, prompt_len:], skip_special_tokens=True)) | |
| def _plain(text: str) -> str: | |
| return text.replace("#", "").replace("*", "").strip() | |
| def _format_block(label: str, text: str) -> str: | |
| return format_block(label, text) | |
| def _detect(image: Image.Image): | |
| return detect_layout(layout_model, image) | |
| def _progress_view(step: int, total: int | None) -> str: | |
| if not total or step <= 0: | |
| percent = 0 | |
| label = "0%" | |
| else: | |
| percent = min(100, round(100 * step / total)) | |
| label = f"{percent}%" | |
| return ( | |
| "<div style=\"margin:4px 0 12px\">" | |
| f"<div style=\"font-size:14px;margin-bottom:6px\">{label}</div>" | |
| "<div style=\"height:10px;background:#e5e7eb;border-radius:999px;overflow:hidden\">" | |
| f"<div style=\"width:{percent}%;height:100%;background:#2563eb\"></div>" | |
| "</div></div>" | |
| ) | |
| def _jobs(image: Image.Image, recognize_seal: bool, recognize_chart: bool): | |
| skipped = set(SKIP_VLM) | |
| if not recognize_seal: | |
| skipped.add("seal") | |
| if not recognize_chart: | |
| skipped.add("chart") | |
| pieces = [] | |
| jobs = [] | |
| for label, box in _detect(image): | |
| if label in skipped: | |
| if label == "image": | |
| pieces.append(_format_block(label, "")) | |
| continue | |
| crop = _crop(image, box) | |
| if crop is None: | |
| continue | |
| jobs.append((label, crop)) | |
| if not jobs and not pieces: | |
| jobs.append(("text", image)) | |
| return pieces, jobs | |
| def parse(image, upright, recognize_seal, recognize_chart, language): | |
| text = _copy(language) | |
| page = image.convert("RGB") | |
| if upright: | |
| page = _upright(page) | |
| pieces, jobs = _jobs(page, bool(recognize_seal), bool(recognize_chart)) | |
| total = len(jobs) + 1 | |
| yield _progress_view(1, total), text["parsing"] | |
| for index, (label, crop) in enumerate(jobs, start=1): | |
| prompt, limit = PROMPTS.get(label, ("Text Recognition:", 1024)) | |
| recognized = _recognize(crop, prompt, limit) | |
| if recognized and label not in SKIP_MARKDOWN: | |
| formatted = _format_block(label, recognized) | |
| if not pieces or _plain(formatted) != _plain(pieces[-1]): | |
| pieces.append(formatted) | |
| yield _progress_view(1 + index, total), text["parsing"] | |
| yield _progress_view(total, total), "\n\n".join(pieces) or text["empty"] | |
| def run(image, upright, recognize_seal, recognize_chart, language): | |
| text = _copy(language) | |
| if image is None: | |
| raise gr.Error(text["need_image"]) | |
| notice = text["parsing"] | |
| yield _progress_view(0, None), notice, notice | |
| for progress, body in parse(image, bool(upright), bool(recognize_seal), bool(recognize_chart), language): | |
| yield progress, body, body | |
| if os.environ.get("OCR_SKIP_LAUNCH") != "1": | |
| import gradio as gr | |
| english = COPY["en"] | |
| def apply_language(language, view): | |
| text = _copy(language) | |
| return ( | |
| gr.update(label=text["language_label"]), | |
| gr.update(value=text["intro"]), | |
| gr.update(value=_note(language)), | |
| gr.update(label=text["image"], placeholder=text["upload"]), | |
| gr.update(label=text["upright"], info=text["upright_info"]), | |
| gr.update(label=text["seal"], info=text["seal_info"]), | |
| gr.update(label=text["chart"], info=text["chart_info"]), | |
| gr.update(value=text["parse"]), | |
| gr.update( | |
| label=text["view_label"], | |
| choices=[(text["rendered"], "rendered"), (text["source"], "source")], | |
| value=view, | |
| ), | |
| ) | |
| def apply_view(view): | |
| rendered = view != "source" | |
| return gr.update(visible=rendered), gr.update(visible=not rendered) | |
| with gr.Blocks(title="MiniCPM-OCR") as demo: | |
| intro = gr.Markdown(english["intro"]) | |
| note = gr.HTML(_note("en")) | |
| with gr.Row(elem_classes=["ocr-columns"]): | |
| with gr.Column(elem_classes=["ocr-col"], scale=1, min_width=0): | |
| language = gr.Radio( | |
| choices=[("English", "en"), ("中文", "zh")], | |
| value="en", | |
| label=english["language_label"], | |
| ) | |
| image = gr.Image(type="pil", label=english["image"], placeholder=english["upload"]) | |
| with gr.Group(): | |
| upright = gr.Checkbox(label=english["upright"], info=english["upright_info"], value=True) | |
| recognize_seal = gr.Checkbox(label=english["seal"], info=english["seal_info"], value=True) | |
| recognize_chart = gr.Checkbox(label=english["chart"], info=english["chart_info"], value=True) | |
| button = gr.Button(english["parse"], variant="primary") | |
| with gr.Column(elem_classes=["ocr-col"], scale=1, min_width=0): | |
| view = gr.Radio( | |
| choices=[(english["rendered"], "rendered"), (english["source"], "source")], | |
| value="rendered", | |
| label=english["view_label"], | |
| ) | |
| progress = gr.HTML() | |
| markdown = gr.Markdown( | |
| elem_classes=["ocr-result"], | |
| sanitize_html=False, | |
| latex_delimiters=[ | |
| {"left": "$$", "right": "$$", "display": True}, | |
| {"left": "$", "right": "$", "display": False}, | |
| ], | |
| max_height=640, | |
| ) | |
| source = gr.Code(language="markdown", visible=False, lines=24, show_label=False) | |
| language.change( | |
| apply_language, | |
| inputs=[language, view], | |
| outputs=[language, intro, note, image, upright, recognize_seal, recognize_chart, button, view], | |
| ) | |
| view.change(apply_view, inputs=[view], outputs=[markdown, source]) | |
| button.click( | |
| run, | |
| inputs=[image, upright, recognize_seal, recognize_chart, language], | |
| outputs=[progress, markdown, source], | |
| show_progress="full", | |
| ) | |
| demo.launch(css=RESULT_CSS) | |