Spaces:
Running on Zero
Running on Zero
| """ | |
| PaddleOCR-VL Doc Editor | |
| ======================= | |
| Interactive document-parsing workbench: | |
| PDF / image -> layout zones (PP-DocLayout, CPU) | |
| -> per-zone recognition (PaddleOCR-VL-0.9B VLM, ZeroGPU) | |
| -> click a zone to edit its text | |
| -> rebuild Markdown / JSON | |
| The two ML stacks are deliberately kept apart: | |
| * Layout detection uses PaddlePaddle on CPU (no CUDA -> no conflict with the | |
| ZeroGPU torch hijack). | |
| * The VLM recognizer uses transformers on the ZeroGPU GPU. | |
| """ | |
| import spaces # must be imported before torch / any CUDA-touching import | |
| import io | |
| import json | |
| import os | |
| import tempfile | |
| import traceback | |
| import fitz # PyMuPDF | |
| import numpy as np | |
| import torch | |
| import gradio as gr | |
| from PIL import Image, ImageDraw, ImageFont | |
| # --------------------------------------------------------------------------- | |
| # Config | |
| # --------------------------------------------------------------------------- | |
| MODEL_PATH = "PaddlePaddle/PaddleOCR-VL" | |
| # Standalone layout model exposed through paddleocr.LayoutDetection. | |
| LAYOUT_MODELS = ["PP-DocLayoutV2", "PP-DocLayout_plus-L", "PP-DocLayout-L"] | |
| # Text-line detector (for dense OCR of scattered text / numbers in images). | |
| TEXTDET_MODELS = ["PP-OCRv5_mobile_det", "PP-OCRv5_server_det", "PP-OCRv4_mobile_det"] | |
| MAX_DENSE_ZONES = 120 # cap text-line zones per page to keep GPU time bounded | |
| RENDER_DPI = 150 # PDF rasterization DPI | |
| MAX_SIDE = 2200 # cap the long side of a page (px) | |
| SAMPLE_PATH = "examples/paddleocr_vl_demo.png" | |
| # On ZeroGPU the GPU is not attached at import time, so torch.cuda.is_available() | |
| # is False here; the `spaces` package intercepts `.to("cuda")` at module scope and | |
| # streams the weights in on the first @spaces.GPU call. So trust the ZeroGPU env | |
| # var, not is_available(), when deciding to place weights on cuda. | |
| ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU")) | |
| USE_CUDA = ON_ZEROGPU or torch.cuda.is_available() | |
| DEVICE = "cuda" if USE_CUDA else "cpu" | |
| # Prompt per PaddleOCR-VL task. | |
| PROMPTS = { | |
| "ocr": "OCR:", | |
| "table": "Table Recognition:", | |
| "formula": "Formula Recognition:", | |
| "chart": "Chart Recognition:", | |
| } | |
| # Max new tokens per task (tables/charts can be long). | |
| MAX_TOKENS = {"ocr": 1024, "table": 1536, "formula": 512, "chart": 1024} | |
| # Free ZeroGPU caps per-call duration (the scheduler scales the declared duration | |
| # up by a safety factor, so keep it modest) AND caps total GPU-seconds/day per | |
| # visitor. Process zones in chunks: few enough allocations to be quota-friendly, | |
| # short enough per call to stay under the duration cap. | |
| GPU_DURATION = 55 | |
| CHUNK = 12 | |
| # Map a layout label to a recognition task. Anything not listed -> "ocr". | |
| LABEL_TO_TASK = { | |
| "table": "table", | |
| "formula": "formula", | |
| "equation": "formula", | |
| "chart": "chart", | |
| "figure": "figure", # image region, skip text recognition | |
| "image": "figure", | |
| "seal": "figure", | |
| "stamp": "figure", | |
| } | |
| # Distinct-ish colors for zone overlays. | |
| PALETTE = [ | |
| (239, 68, 68), (59, 130, 246), (16, 185, 129), (245, 158, 11), | |
| (139, 92, 246), (236, 72, 153), (14, 165, 233), (34, 197, 94), | |
| (249, 115, 22), (168, 85, 247), (20, 184, 166), (99, 102, 241), | |
| ] | |
| # --------------------------------------------------------------------------- | |
| # Models (lazy singletons) | |
| # --------------------------------------------------------------------------- | |
| _vlm = {"model": None, "processor": None} | |
| _layout = {"model": None, "name": None} | |
| _textdet = {"model": None, "name": None} | |
| def _patch_transformers_masking(): | |
| """Bridge a kwarg-name mismatch in the model's remote code. | |
| PaddleOCR-VL's modeling code calls | |
| create_causal_mask(inputs_embeds=...) | |
| but transformers 4.x names that parameter `input_embeds` (no 's'). Wrap the | |
| function so the remote code loads and runs unchanged. Must run before the | |
| remote modeling module does `from transformers.masking_utils import ...`. | |
| """ | |
| import functools | |
| try: | |
| import transformers.masking_utils as mu | |
| except Exception as e: # noqa: BLE001 | |
| print("[patch] masking_utils unavailable:", e, flush=True) | |
| return | |
| orig = mu.create_causal_mask | |
| if getattr(orig, "_paddleocr_shim", False): | |
| return | |
| def shim(*args, **kwargs): | |
| if "inputs_embeds" in kwargs and "input_embeds" not in kwargs: | |
| kwargs["input_embeds"] = kwargs.pop("inputs_embeds") | |
| return orig(*args, **kwargs) | |
| shim._paddleocr_shim = True | |
| mu.create_causal_mask = shim | |
| print("[patch] create_causal_mask kwarg shim installed", flush=True) | |
| def _load_vlm(): | |
| """Load the PaddleOCR-VL VLM at module import so ZeroGPU can pack the weights.""" | |
| if _vlm["model"] is not None: | |
| return | |
| _patch_transformers_masking() | |
| from transformers import AutoModelForCausalLM, AutoProcessor | |
| print("[vlm] loading", MODEL_PATH, flush=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_PATH, | |
| trust_remote_code=True, | |
| torch_dtype=torch.bfloat16, | |
| ).eval() | |
| if DEVICE == "cuda": | |
| model = model.to("cuda") | |
| processor = AutoProcessor.from_pretrained(MODEL_PATH, trust_remote_code=True) | |
| _vlm["model"] = model | |
| _vlm["processor"] = processor | |
| print("[vlm] ready", flush=True) | |
| def _get_layout(): | |
| """Lazily build the CPU layout detector, trying model names in order.""" | |
| if _layout["model"] is not None: | |
| return _layout["model"] | |
| from paddleocr import LayoutDetection | |
| last_err = None | |
| for name in LAYOUT_MODELS: | |
| try: | |
| print("[layout] trying", name, flush=True) | |
| _layout["model"] = LayoutDetection(model_name=name) | |
| _layout["name"] = name | |
| print("[layout] using", name, flush=True) | |
| return _layout["model"] | |
| except Exception as e: # noqa: BLE001 | |
| last_err = e | |
| print(f"[layout] {name} failed: {e}", flush=True) | |
| raise RuntimeError(f"Could not init any layout model: {last_err}") | |
| def _get_textdet(): | |
| """Lazily build the CPU text-line detector (PP-OCR detection).""" | |
| if _textdet["model"] is not None: | |
| return _textdet["model"] | |
| from paddleocr import TextDetection | |
| last_err = None | |
| for name in TEXTDET_MODELS: | |
| try: | |
| print("[textdet] trying", name, flush=True) | |
| _textdet["model"] = TextDetection(model_name=name) | |
| _textdet["name"] = name | |
| print("[textdet] using", name, flush=True) | |
| return _textdet["model"] | |
| except Exception as e: # noqa: BLE001 | |
| last_err = e | |
| print(f"[textdet] {name} failed: {e}", flush=True) | |
| raise RuntimeError(f"Could not init any text-detection model: {last_err}") | |
| # Load the VLM eagerly at import (ZeroGPU best practice). | |
| try: | |
| _load_vlm() | |
| except Exception as e: # noqa: BLE001 | |
| print("[vlm] eager load failed, will retry lazily:", e, flush=True) | |
| # --------------------------------------------------------------------------- | |
| # PDF / image helpers | |
| # --------------------------------------------------------------------------- | |
| def _cap_size(img: Image.Image) -> Image.Image: | |
| w, h = img.size | |
| m = max(w, h) | |
| if m > MAX_SIDE: | |
| s = MAX_SIDE / m | |
| img = img.resize((int(w * s), int(h * s)), Image.LANCZOS) | |
| return img | |
| def load_pages(file_path): | |
| """Return a list of PIL.Image pages from a PDF or single image file.""" | |
| if not file_path: | |
| return [] | |
| ext = os.path.splitext(file_path)[1].lower() | |
| pages = [] | |
| if ext == ".pdf": | |
| doc = fitz.open(file_path) | |
| zoom = RENDER_DPI / 72.0 | |
| mat = fitz.Matrix(zoom, zoom) | |
| for page in doc: | |
| pix = page.get_pixmap(matrix=mat, alpha=False) | |
| img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB") | |
| pages.append(_cap_size(img)) | |
| doc.close() | |
| else: | |
| pages.append(_cap_size(Image.open(file_path).convert("RGB"))) | |
| return pages | |
| # --------------------------------------------------------------------------- | |
| # Layout detection (CPU) | |
| # --------------------------------------------------------------------------- | |
| def _finalize(zones): | |
| """Sort into rough reading order and assign ids.""" | |
| zones.sort(key=lambda z: (z["box"][1], z["box"][0])) | |
| for i, z in enumerate(zones): | |
| z["id"] = i | |
| return zones | |
| def detect_layout_zones(img: Image.Image): | |
| """PP-DocLayout: region-level zones (text, title, table, formula, chart, figure).""" | |
| model = _get_layout() | |
| arr = np.array(img) | |
| out = model.predict(arr, batch_size=1) | |
| res = list(out)[0] if not isinstance(out, dict) else out | |
| boxes = res["boxes"] if "boxes" in res else res.get("boxes", []) | |
| zones = [] | |
| W, H = img.size | |
| for b in boxes: | |
| coord = b.get("coordinate") or b.get("bbox") or b.get("box") | |
| if coord is None: | |
| continue | |
| x1, y1, x2, y2 = [float(v) for v in coord[:4]] | |
| x1, y1 = max(0, x1), max(0, y1) | |
| x2, y2 = min(W, x2), min(H, y2) | |
| if x2 - x1 < 3 or y2 - y1 < 3: | |
| continue | |
| label = str(b.get("label", "text")).lower().strip() | |
| task = LABEL_TO_TASK.get(label, "ocr") | |
| zones.append( | |
| { | |
| "box": [int(x1), int(y1), int(x2), int(y2)], | |
| "label": label, | |
| "task": task, | |
| "score": float(b.get("score", 0.0)), | |
| } | |
| ) | |
| return _finalize(zones) | |
| def detect_text_line_zones(img: Image.Image): | |
| """PP-OCR text detection: one zone per detected text line / number.""" | |
| model = _get_textdet() | |
| arr = np.array(img) | |
| out = model.predict(arr, batch_size=1) | |
| res = list(out)[0] if not isinstance(out, dict) else out | |
| polys = res.get("dt_polys", res.get("polys", [])) | |
| scores = res.get("dt_scores", res.get("scores", [])) | |
| zones = [] | |
| W, H = img.size | |
| for i, p in enumerate(polys): | |
| pts = np.array(p).reshape(-1, 2) | |
| x1, y1 = pts[:, 0].min(), pts[:, 1].min() | |
| x2, y2 = pts[:, 0].max(), pts[:, 1].max() | |
| x1, y1 = max(0, float(x1)), max(0, float(y1)) | |
| x2, y2 = min(W, float(x2)), min(H, float(y2)) | |
| bw, bh = x2 - x1, y2 - y1 | |
| if bw < 2 or bh < 2: | |
| continue | |
| # A much taller-than-wide box is vertical text -> keep it vertical. | |
| vertical = bh > bw * 1.8 | |
| angle = 90 if vertical else 0 | |
| # pad along the reading direction to help the recognizer | |
| if vertical: | |
| pad = int(bw * 0.15) | |
| box = [max(0, int(x1 - pad)), int(y1), min(W, int(x2 + pad)), int(y2)] | |
| else: | |
| pad = int(bh * 0.15) | |
| box = [int(x1), max(0, int(y1 - pad)), int(x2), min(H, int(y2 + pad))] | |
| zones.append( | |
| { | |
| "box": box, | |
| "label": "text-line", | |
| "task": "ocr", | |
| "score": float(scores[i]) if i < len(scores) else 0.0, | |
| "dense": True, | |
| "angle": angle, | |
| } | |
| ) | |
| if len(zones) > MAX_DENSE_ZONES: | |
| # keep the highest-confidence lines | |
| zones.sort(key=lambda z: z["score"], reverse=True) | |
| zones = zones[:MAX_DENSE_ZONES] | |
| return _finalize(zones) | |
| def detect_zones(img: Image.Image, mode: str = "Document layout"): | |
| if mode.startswith("Dense"): | |
| return detect_text_line_zones(img) | |
| return detect_layout_zones(img) | |
| # --------------------------------------------------------------------------- | |
| # VLM recognition (GPU) | |
| # --------------------------------------------------------------------------- | |
| def _vlm_one(crop: Image.Image, task: str, max_tokens: int) -> str: | |
| model = _vlm["model"] | |
| processor = _vlm["processor"] | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": crop}, | |
| {"type": "text", "text": PROMPTS[task]}, | |
| ], | |
| } | |
| ] | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| tokenize=True, | |
| add_generation_prompt=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ).to(model.device) | |
| with torch.inference_mode(): | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=max_tokens, | |
| do_sample=False, | |
| use_cache=True, | |
| ) | |
| # Decode only the newly generated tokens, not the prompt echo. | |
| in_len = inputs["input_ids"].shape[1] | |
| gen = out[:, in_len:] | |
| text = processor.batch_decode(gen, skip_special_tokens=True)[0] | |
| # Belt-and-suspenders: if a chat template still leaks through, keep the tail. | |
| for marker in ("Assistant:", "assistant\n", "assistant:"): | |
| if marker in text: | |
| text = text.split(marker)[-1] | |
| return text.strip() | |
| def _recognize_chunk(img: Image.Image, chunk, ocr_figures: bool): | |
| """Recognize a small batch of zones inside one short GPU allocation. | |
| Runs in the ZeroGPU worker process, so results are returned (not mutated in | |
| place) to cross the process boundary. | |
| """ | |
| _load_vlm() | |
| for z in chunk: | |
| task = z["task"] | |
| if task == "figure": | |
| if not ocr_figures: | |
| z["content"] = "" | |
| continue | |
| task = "ocr" # OCR the figure region instead of skipping it | |
| # Dense single-line boxes need few tokens; regions can need many. | |
| max_tokens = 256 if z.get("dense") else MAX_TOKENS.get(task, 1024) | |
| crop = img.crop(tuple(z["box"])) | |
| # Rotate a vertical crop to horizontal so the VLM reads it correctly; the | |
| # field/text is placed back at the same angle when the PDF is built. | |
| if int(z.get("angle", 0) or 0) == 90: | |
| crop = crop.transpose(Image.ROTATE_90) | |
| try: | |
| z["content"] = _vlm_one(crop, task, max_tokens) | |
| except Exception as e: # noqa: BLE001 | |
| z["content"] = f"[recognition error: {e}]" | |
| return chunk | |
| def recognize_all(img: Image.Image, zones, ocr_figures: bool = True): | |
| """Recognize all zones, chunked across several short GPU allocations. | |
| ocr_figures: if True, image/figure regions are also OCR'd (pulls out text / | |
| numbers embedded in diagrams); if False they are left blank. | |
| """ | |
| for i in range(0, len(zones), CHUNK): | |
| done = _recognize_chunk(img, zones[i : i + CHUNK], ocr_figures) | |
| zones[i : i + CHUNK] = done | |
| return zones | |
| # --------------------------------------------------------------------------- | |
| # Rendering helpers | |
| # --------------------------------------------------------------------------- | |
| def _font(size=22): | |
| try: | |
| return ImageFont.truetype("DejaVuSans-Bold.ttf", size) | |
| except Exception: # noqa: BLE001 | |
| return ImageFont.load_default() | |
| def render_base(img: Image.Image, zones): | |
| """Draw numbered rectangles on a copy of the page for display.""" | |
| base = img.convert("RGB").copy() | |
| draw = ImageDraw.Draw(base) | |
| f = _font(max(16, img.size[0] // 60)) | |
| for z in zones: | |
| color = PALETTE[z["id"] % len(PALETTE)] | |
| x1, y1, x2, y2 = z["box"] | |
| draw.rectangle([x1, y1, x2, y2], outline=color, width=3) | |
| tag = str(z["id"] + 1) | |
| tw = draw.textlength(tag, font=f) | |
| draw.rectangle([x1, y1, x1 + tw + 8, y1 + 26], fill=color) | |
| draw.text((x1 + 4, y1 + 2), tag, fill=(255, 255, 255), font=f) | |
| return base | |
| def build_annotations(img: Image.Image, zones): | |
| """AnnotatedImage value: (base_image, [(mask, label), ...]).""" | |
| base = render_base(img, zones) | |
| W, H = img.size | |
| anns = [] | |
| for z in zones: | |
| mask = np.zeros((H, W), dtype=np.uint8) | |
| x1, y1, x2, y2 = z["box"] | |
| mask[y1:y2, x1:x2] = 1 | |
| anns.append((mask, f"{z['id'] + 1}. {z['label']}")) | |
| return (base, anns) | |
| def zone_choices(zones): | |
| return [(f"{z['id'] + 1}. {z['label']}", z["id"]) for z in zones] | |
| # --------------------------------------------------------------------------- | |
| # Document rebuild | |
| # --------------------------------------------------------------------------- | |
| def rebuild_markdown(zones): | |
| ordered = sorted(zones, key=lambda z: (z["box"][1], z["box"][0])) | |
| parts = [] | |
| for z in ordered: | |
| c = (z.get("content") or "").strip() | |
| if z["task"] == "figure" and not c: | |
| parts.append(f"*[figure — zone {z['id'] + 1}]*") | |
| continue | |
| if not c: | |
| continue | |
| if z["label"] in ("title", "doc_title", "paragraph_title"): | |
| parts.append(f"## {c}") | |
| elif z["task"] == "formula": | |
| parts.append(f"$$\n{c}\n$$") | |
| else: | |
| parts.append(c) | |
| return "\n\n".join(parts) | |
| # --------------------------------------------------------------------------- | |
| # PDF rebuild | |
| # --------------------------------------------------------------------------- | |
| def _zones_for(zones_pages, i): | |
| zp = zones_pages or {} | |
| return zp.get(i) or zp.get(str(i)) or [] | |
| def _add_form_field(page, z, text, page_index): | |
| """Add an editable AcroForm text field at the zone, pre-filled with `text`.""" | |
| x1, y1, x2, y2 = z["box"] | |
| rect = fitz.Rect(x1, y1, x2, y2) | |
| # Paint an opaque white rectangle first to hide the original scanned text, so | |
| # the editable field's text is not shown on top of a duplicate underneath. | |
| page.draw_rect(rect, fill=(1, 1, 1), color=(0.7, 0.75, 0.85), width=0.5) | |
| w = fitz.Widget() | |
| w.rect = rect | |
| # Unique across the whole document: fields sharing a name are one field in | |
| # PDF editors, so editing page 1 would also change page 2. | |
| w.field_name = f"p{page_index + 1}_zone_{z['id'] + 1}" | |
| w.field_type = fitz.PDF_WIDGET_TYPE_TEXT | |
| w.field_flags = fitz.PDF_TX_FIELD_IS_MULTILINE | |
| w.field_value = text | |
| w.text_fontsize = 0 # auto-size to the box | |
| w.text_font = "Helv" | |
| # Transparent field (the white rectangle above already provides the clean | |
| # background); no border so the box reads as plain white. | |
| w.fill_color = None | |
| w.border_color = None | |
| annot = page.add_widget(w) | |
| angle = int(z.get("angle", 0) or 0) % 360 | |
| if angle: | |
| # Keep the original orientation (vertical labels stay vertical). PDF | |
| # field rotation lives in /MK /R; PyMuPDF's Widget has no rotation | |
| # setting, so set the key and regenerate the appearance. | |
| doc = page.parent | |
| doc.xref_set_key(annot.xref, "MK/R", str(angle)) | |
| placed = next((x for x in page.widgets() if x.xref == annot.xref), None) | |
| if placed is not None: | |
| placed.update() | |
| if "/R" not in (doc.xref_get_key(annot.xref, "MK")[1] or ""): | |
| doc.xref_set_key(annot.xref, "MK/R", str(angle)) | |
| def _finalize_form(doc): | |
| """Make the fields editable in any PDF editor. | |
| Editors redraw a field from its /DA font when you type and resolve that font | |
| through the form's /DR resources; declare it, and ask viewers to regenerate | |
| appearances so values always render. | |
| """ | |
| catalog = doc.pdf_catalog() | |
| if doc.xref_get_key(catalog, "AcroForm/DR/Font/Helv")[0] == "null": | |
| xref = doc.get_new_xref() | |
| doc.update_object( | |
| xref, "<</Type/Font/Subtype/Type1/BaseFont/Helvetica/Encoding/WinAnsiEncoding>>" | |
| ) | |
| doc.xref_set_key(catalog, "AcroForm/DR/Font/Helv", f"{xref} 0 R") | |
| if doc.xref_get_key(catalog, "AcroForm/DA")[0] == "null": | |
| doc.xref_set_key(catalog, "AcroForm/DA", "(/Helv 0 Tf 0 g)") | |
| try: | |
| doc.need_appearances(True) | |
| except Exception: # noqa: BLE001 | |
| pass | |
| def build_pdf(pages, zones_pages, style="Searchable (over original)"): | |
| """Rebuild a PDF from the pages + recognized/edited zones. | |
| Searchable: original page image as background with an invisible, selectable | |
| text layer placed at each zone (looks identical to the input, but the text is | |
| now real and reflects your edits). | |
| Editable form: original page image with a fillable AcroForm text field at each | |
| zone, pre-filled with the recognized text — edit it in any PDF viewer. | |
| Text only: white pages with the recognized text drawn at each zone's position. | |
| """ | |
| searchable = style.startswith("Searchable") | |
| form = style.startswith("Editable") | |
| doc = fitz.open() | |
| for i, img in enumerate(pages): | |
| W, H = img.size | |
| page = doc.new_page(width=W, height=H) | |
| if searchable or form: | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| page.insert_image(fitz.Rect(0, 0, W, H), stream=buf.getvalue()) | |
| for z in _zones_for(zones_pages, i): | |
| text = (z.get("content") or "").strip() | |
| if not text: | |
| continue | |
| if form: | |
| try: | |
| _add_form_field(page, z, text, i) | |
| except Exception: # noqa: BLE001 | |
| traceback.print_exc() | |
| continue | |
| x1, y1, x2, y2 = z["box"] | |
| rect = fitz.Rect(x1, y1, x2, y2) | |
| bh = max(1, y2 - y1) | |
| if z.get("dense"): | |
| fs = max(6, min(int(bh * 0.7), 40)) | |
| else: | |
| fs = 10 | |
| render_mode = 3 if searchable else 0 # 3 = invisible (search layer) | |
| rotate = int(z.get("angle", 0) or 0) | |
| for fname in ("china-s", None): | |
| try: | |
| kw = {"fontname": fname} if fname else {} | |
| page.insert_textbox( | |
| rect, text, fontsize=fs, render_mode=render_mode, | |
| align=0, rotate=rotate, **kw, | |
| ) | |
| break | |
| except Exception: # noqa: BLE001 | |
| continue | |
| if form and any(_zones_for(zones_pages, i) for i in range(len(pages))): | |
| _finalize_form(doc) | |
| out = os.path.join(tempfile.mkdtemp(), "document.pdf") | |
| doc.save(out, deflate=True, garbage=3) | |
| doc.close() | |
| return out | |
| # --------------------------------------------------------------------------- | |
| # Gradio callbacks | |
| # --------------------------------------------------------------------------- | |
| def on_upload(file_path): | |
| pages = load_pages(file_path) | |
| if not pages: | |
| return ( | |
| gr.update(choices=[], value=None), | |
| [], | |
| {}, | |
| [], | |
| None, | |
| gr.update(choices=[], value=None), | |
| "", | |
| "", | |
| "Upload a PDF or image to begin.", | |
| ) | |
| choices = [(f"Page {i + 1}", i) for i in range(len(pages))] | |
| first = pages[0] | |
| return ( | |
| gr.update(choices=choices, value=0), | |
| pages, | |
| {}, | |
| [], | |
| (first, []), | |
| gr.update(choices=[], value=None), | |
| "", | |
| "", | |
| f"Loaded {len(pages)} page(s). Click **Process pages & export PDF** to do them all, " | |
| "or pick a page and click **Detect & recognize this page**.", | |
| ) | |
| def load_sample(): | |
| if not os.path.exists(SAMPLE_PATH): | |
| return on_upload(None) | |
| return on_upload(SAMPLE_PATH) | |
| def on_page_change(page_idx, pages, zones_pages): | |
| if not pages or page_idx is None or int(page_idx) >= len(pages): | |
| return (None, [], gr.update(choices=[], value=None), "", "") | |
| i = int(page_idx) | |
| img = pages[i] | |
| zones = _zones_for(zones_pages, i) | |
| if zones: | |
| return ( | |
| build_annotations(img, zones), | |
| zones, | |
| gr.update(choices=zone_choices(zones), value=zones[0]["id"]), | |
| zones[0]["label"], | |
| zones[0].get("content", ""), | |
| ) | |
| return ((img, []), [], gr.update(choices=[], value=None), "", "") | |
| def _ensure_pages(pages, file_path): | |
| """Pages from session state, or straight from the uploaded file. | |
| Clicking a button right after uploading can run before the upload handler has | |
| stored the pages, so fall back to reading the file itself. | |
| """ | |
| if pages: | |
| return pages | |
| if file_path: | |
| loaded = load_pages(file_path) | |
| if loaded: | |
| return loaded | |
| raise gr.Error("Upload a document (or click 'Load sample') first.") | |
| def _page_choices(pages, page_idx): | |
| return gr.update(choices=[(f"Page {i + 1}", i) for i in range(len(pages))], value=page_idx) | |
| def analyze(file_path, page_idx, pages, mode, ocr_figures, zones_pages): | |
| pages = _ensure_pages(pages, file_path) | |
| if page_idx is None or int(page_idx) >= len(pages): | |
| page_idx = 0 | |
| page_idx = int(page_idx) | |
| img = pages[page_idx] | |
| try: | |
| zones = detect_zones(img, mode) | |
| except Exception as e: # noqa: BLE001 | |
| traceback.print_exc() | |
| raise gr.Error(f"Detection failed: {e}") | |
| zp = dict(zones_pages or {}) | |
| if not zones: | |
| zp[page_idx] = [] | |
| return ( | |
| pages, | |
| _page_choices(pages, page_idx), | |
| (img, []), | |
| [], | |
| zp, | |
| gr.update(choices=[], value=None), | |
| "", | |
| "", | |
| "No zones detected on this page. Try the other detection mode.", | |
| ) | |
| zones = recognize_all(img, zones, bool(ocr_figures)) | |
| zp[page_idx] = zones | |
| ann = build_annotations(img, zones) | |
| status = ( | |
| f"Recognized {len(zones)} zone(s) on page {page_idx + 1}. Click a box to " | |
| "edit, then **Export PDF**." | |
| ) | |
| return ( | |
| pages, | |
| _page_choices(pages, page_idx), | |
| ann, | |
| zones, | |
| zp, | |
| gr.update(choices=zone_choices(zones), value=zones[0]["id"]), | |
| zones[0]["label"], | |
| zones[0].get("content", ""), | |
| status, | |
| ) | |
| def _select_zone(zones, zid): | |
| for z in zones: | |
| if z["id"] == zid: | |
| return z["label"], z.get("content", "") | |
| return "", "" | |
| def on_box_select(zones, evt: gr.SelectData): | |
| if not zones: | |
| return gr.update(), "", "" | |
| idx = evt.index | |
| if isinstance(idx, (list, tuple)): | |
| idx = idx[0] | |
| # annotations are built in zone-id order | |
| if idx is None or idx < 0 or idx >= len(zones): | |
| return gr.update(), "", "" | |
| z = zones[idx] | |
| label, content = _select_zone(zones, z["id"]) | |
| return gr.update(value=z["id"]), label, content | |
| def on_dropdown(zones, zid): | |
| if zid is None or not zones: | |
| return "", "" | |
| return _select_zone(zones, zid) | |
| def save_edit(zones, zid, text, page_idx, zones_pages): | |
| if zid is None or not zones: | |
| return zones, zones_pages, "Select a zone first." | |
| for z in zones: | |
| if z["id"] == zid: | |
| z["content"] = text | |
| break | |
| else: | |
| return zones, zones_pages, "Zone not found." | |
| zp = dict(zones_pages or {}) | |
| if page_idx is not None: | |
| zp[int(page_idx)] = zones | |
| return zones, zp, f"Saved zone {zid + 1}. Click **Export PDF** when done." | |
| def api_parse_sample(mode: str = "Document layout", ocr_figures: bool = True): | |
| """API-only: parse the bundled sample end to end and return the result. | |
| Lets the pipeline be smoke-tested (and called via API/MCP) without a file | |
| upload. Returns detected zones with their recognized content and the rebuilt | |
| Markdown. | |
| """ | |
| pages = load_pages(SAMPLE_PATH) | |
| if not pages: | |
| return {"error": "sample not found"} | |
| img = pages[0] | |
| zones = detect_zones(img, mode) | |
| zones = recognize_all(img, zones, bool(ocr_figures)) | |
| return { | |
| "mode": mode, | |
| "n_zones": len(zones), | |
| "markdown": rebuild_markdown(zones), | |
| "zones": [ | |
| {"id": z["id"], "label": z["label"], "box": z["box"], | |
| "content": z.get("content", "")} | |
| for z in zones | |
| ], | |
| } | |
| def export_pdf(pages, zones_pages, style): | |
| if not pages: | |
| return None, "Upload a document first." | |
| if not any(_zones_for(zones_pages, i) for i in range(len(pages))): | |
| return None, "No recognized zones yet — click **Process pages & export PDF** first." | |
| try: | |
| path = build_pdf(pages, zones_pages, style) | |
| except Exception as e: # noqa: BLE001 | |
| traceback.print_exc() | |
| raise gr.Error(f"PDF export failed: {e}") | |
| n = sum(1 for i in range(len(pages)) if _zones_for(zones_pages, i)) | |
| return path, f"PDF ready ({n} recognized page(s)). Download below. ⬇️" | |
| def parse_page_spec(spec: str, page_count: int) -> list[int]: | |
| """'all' / '' -> every page; '1-5, 8, 10-12' -> those pages (0-based, sorted).""" | |
| spec = (spec or "").strip().lower() | |
| if spec in ("", "all", "*"): | |
| return list(range(page_count)) | |
| picked: set[int] = set() | |
| for part in spec.replace(";", ",").split(","): | |
| part = part.strip() | |
| if not part: | |
| continue | |
| try: | |
| if "-" in part: | |
| a, b = part.split("-", 1) | |
| start, end = int(a), int(b) | |
| if start > end: | |
| start, end = end, start | |
| else: | |
| start = end = int(part) | |
| except ValueError: | |
| raise gr.Error(f"Could not read the page list near '{part}'. Use e.g. all, 1-5, 8") | |
| picked.update(range(max(1, start), min(page_count, end) + 1)) | |
| if not picked: | |
| raise gr.Error(f"No pages in range — this document has {page_count} page(s).") | |
| return sorted(p - 1 for p in picked) | |
| def _is_quota_error(exc: Exception) -> bool: | |
| msg = str(exc).lower() | |
| return "quota" in msg or ("zerogpu" in msg and "limit" in msg) | |
| def analyze_pages(file_path, page_spec, pages, mode, ocr_figures, zones_pages, style, page_idx, | |
| progress=gr.Progress()): | |
| """Detect + recognize several pages in one run, then export one PDF.""" | |
| pages = _ensure_pages(pages, file_path) | |
| if page_idx is None or int(page_idx) >= len(pages): | |
| page_idx = 0 | |
| targets = parse_page_spec(page_spec, len(pages)) | |
| zp = dict(zones_pages or {}) | |
| done, failed, stopped_at = [], [], None | |
| for k, i in enumerate(targets): | |
| progress((k, len(targets)), desc=f"Page {i + 1} ({k + 1} of {len(targets)})") | |
| img = pages[i] | |
| try: | |
| zones = detect_zones(img, mode) | |
| zones = recognize_all(img, zones, bool(ocr_figures)) if zones else [] | |
| except Exception as e: # noqa: BLE001 | |
| traceback.print_exc() | |
| if _is_quota_error(e): | |
| # Out of free GPU time: keep what finished and export that. | |
| stopped_at = i + 1 | |
| break | |
| failed.append(i + 1) | |
| continue | |
| zp[i] = zones | |
| done.append(i + 1) | |
| progress((len(targets), len(targets)), desc="Building PDF") | |
| path = None | |
| if any(_zones_for(zp, i) for i in range(len(pages))): | |
| path = build_pdf(pages, zp, style) | |
| parts = [f"Processed {len(done)} of {len(targets)} page(s)."] | |
| if path: | |
| parts.append("The PDF below has every page; processed pages have editable fields. ⬇️") | |
| if failed: | |
| parts.append(f"Could not read page(s) {', '.join(map(str, failed))} — left unchanged.") | |
| if stopped_at: | |
| parts.append( | |
| f"Stopped at page {stopped_at}: the free GPU quota ran out. Sign in to Hugging " | |
| "Face for more quota, then process the remaining pages — finished pages are kept." | |
| ) | |
| if not path: | |
| parts.append("No text was found to export.") | |
| # Refresh the page currently on screen so its new boxes are clickable. | |
| view = on_page_change(page_idx, pages, zp) | |
| return (pages, _page_choices(pages, int(page_idx)), zp, path, " ".join(parts), *view) | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| LATEX_DELIMS = [ | |
| {"left": "$$", "right": "$$", "display": True}, | |
| {"left": "$", "right": "$", "display": False}, | |
| ] | |
| with gr.Blocks(title="PaddleOCR-VL Doc Editor", theme=gr.themes.Soft()) as demo: | |
| pages_state = gr.State([]) | |
| zones_state = gr.State([]) | |
| zones_pages = gr.State({}) | |
| gr.Markdown( | |
| "# 📝 PaddleOCR-VL Doc Editor\n" | |
| "Upload a **PDF or image** → detected zones are recognized by " | |
| "[PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) → " | |
| "**click any box to edit** its text → **export a corrected PDF**." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| file_in = gr.File( | |
| label="Upload PDF or image", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp"], | |
| type="filepath", | |
| ) | |
| with gr.Row(): | |
| page_sel = gr.Dropdown(label="Page", choices=[], interactive=True, scale=2) | |
| sample_btn = gr.Button("Load sample", scale=1) | |
| analyze_btn = gr.Button("Detect & recognize this page", variant="primary", scale=3) | |
| with gr.Row(): | |
| page_spec = gr.Textbox( | |
| label="Pages to process", | |
| value="all", | |
| placeholder="all, or e.g. 1-5, 8, 10-12", | |
| info="Process several pages in one run and download one PDF.", | |
| scale=2, | |
| ) | |
| batch_btn = gr.Button("⚡ Process pages & export PDF", variant="primary", scale=3) | |
| with gr.Row(): | |
| mode = gr.Radio( | |
| choices=["Document layout", "Dense text lines"], | |
| value="Document layout", | |
| label="Detection mode", | |
| info=( | |
| "Document layout = region zones (text/table/formula/chart) for " | |
| "structured docs. Dense text lines = one box per text/number, " | |
| "best for diagrams, drawings and scattered labels." | |
| ), | |
| scale=3, | |
| ) | |
| ocr_figures = gr.Checkbox( | |
| value=True, | |
| label="OCR image/figure regions too", | |
| scale=1, | |
| ) | |
| annotated = gr.AnnotatedImage( | |
| label="Detected zones (click a box to edit)", | |
| height=640, | |
| ) | |
| with gr.Column(scale=2): | |
| status = gr.Markdown("Upload a PDF or image to begin.") | |
| zone_dd = gr.Dropdown(label="Zone", choices=[], interactive=True) | |
| zone_label = gr.Textbox(label="Zone type", interactive=False) | |
| zone_editor = gr.Textbox(label="Recognized content (editable)", lines=12) | |
| save_btn = gr.Button("💾 Save edit") | |
| pdf_style = gr.Radio( | |
| choices=[ | |
| "Editable form (fillable fields)", | |
| "Searchable (over original)", | |
| "Text only (white pages)", | |
| ], | |
| value="Editable form (fillable fields)", | |
| label="PDF output", | |
| info=( | |
| "Editable form = a fillable field at each zone you edit directly " | |
| "in any PDF viewer. Searchable = original pages with a corrected " | |
| "text layer. Text only = white pages with the recognized text." | |
| ), | |
| ) | |
| export_btn = gr.Button("📄 Export PDF", variant="primary") | |
| dl_pdf = gr.File(label="PDF download") | |
| # --- wiring --- | |
| up_outputs = [ | |
| page_sel, pages_state, zones_pages, zones_state, annotated, | |
| zone_dd, zone_label, zone_editor, status, | |
| ] | |
| file_in.change(on_upload, inputs=file_in, outputs=up_outputs) | |
| sample_btn.click(load_sample, inputs=None, outputs=up_outputs) | |
| page_sel.change( | |
| on_page_change, | |
| inputs=[page_sel, pages_state, zones_pages], | |
| outputs=[annotated, zones_state, zone_dd, zone_label, zone_editor], | |
| ) | |
| analyze_btn.click( | |
| analyze, | |
| inputs=[file_in, page_sel, pages_state, mode, ocr_figures, zones_pages], | |
| outputs=[pages_state, page_sel, annotated, zones_state, zones_pages, zone_dd, zone_label, | |
| zone_editor, status], | |
| ) | |
| annotated.select( | |
| on_box_select, | |
| inputs=[zones_state], | |
| outputs=[zone_dd, zone_label, zone_editor], | |
| ) | |
| zone_dd.change( | |
| on_dropdown, | |
| inputs=[zones_state, zone_dd], | |
| outputs=[zone_label, zone_editor], | |
| ) | |
| save_btn.click( | |
| save_edit, | |
| inputs=[zones_state, zone_dd, zone_editor, page_sel, zones_pages], | |
| outputs=[zones_state, zones_pages, status], | |
| ) | |
| export_btn.click( | |
| export_pdf, | |
| inputs=[pages_state, zones_pages, pdf_style], | |
| outputs=[dl_pdf, status], | |
| ) | |
| batch_btn.click( | |
| analyze_pages, | |
| inputs=[file_in, page_spec, pages_state, mode, ocr_figures, zones_pages, pdf_style, page_sel], | |
| outputs=[pages_state, page_sel, zones_pages, dl_pdf, status, annotated, zones_state, zone_dd, | |
| zone_label, zone_editor], | |
| ) | |
| # Hidden API endpoint: parse the bundled sample (for smoke tests / API use). | |
| with gr.Row(visible=False): | |
| api_btn = gr.Button("api") | |
| api_out = gr.JSON() | |
| api_btn.click( | |
| api_parse_sample, | |
| inputs=[mode, ocr_figures], | |
| outputs=api_out, | |
| api_name="parse_sample", | |
| ) | |
| demo.queue(default_concurrency_limit=1, max_size=16) | |
| if __name__ == "__main__": | |
| # SSR (experimental) proxies requests through a Node server; it intermittently | |
| # returns 502 on file uploads, so serve the app directly. | |
| demo.launch(ssr_mode=False) | |