OCR-Compare / app.py
vivekreddyr's picture
Add click-to-edit PDF export (erase + rewrite, download edited PDF)
6188ea2 verified
Raw History Blame Contribute Delete
22.3 kB
"""
OCR Compare -- classic PaddleOCR vs PaddleOCR-VL, side by side
==============================================================
PDF / image page
├─► classic PaddleOCR (PP-OCRv5 det + rec, CPU) ─┐
└─► PaddleOCR-VL-0.9B (VLM, ZeroGPU) ─┤ same {text, box} shape
▼
two page panels with invisible click-to-edit boxes + a diff table
Everything runs on Hugging Face: PaddlePaddle on the Space CPU, the VLM on
ZeroGPU. The box overlay and comparison live in overlay.py.
The VLM loading code (version pins, kwarg shim, chunked GPU calls, prompt
slicing) is carried over from the PaddleOCR-VL-Doc-Editor Space, where each of
those was needed to make it run on ZeroGPU at all.
"""
try:
import spaces # must be imported before torch / any CUDA-touching import
except ImportError: # local run without the HF `spaces` package
class spaces: # noqa: N801
@staticmethod
def GPU(*args, **kwargs):
if args and callable(args[0]):
return args[0]
return lambda fn: fn
import io
import os
import time
import traceback
import fitz # PyMuPDF
import gradio as gr
import numpy as np
from PIL import Image
from overlay import (
CSS,
HEAD_JS,
TABLE_HEADERS,
compare,
image_to_data_uri,
overlay_html,
placeholder_html,
quad_to_box,
reading_order,
table_rows,
)
from pdf_export import build_edited_pdf, parse_edits
try:
import torch
except ImportError: # local UI testing without the VLM stack
torch = None
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
VL_MODEL = "PaddlePaddle/PaddleOCR-VL"
CLASSIC_MODELS = {
"mobile (fast)": ("PP-OCRv5_mobile_det", "PP-OCRv5_mobile_rec"),
"server (accurate)": ("PP-OCRv5_server_det", "PP-OCRv5_server_rec"),
}
LAYOUT_MODELS = ["PP-DocLayoutV2", "PP-DocLayout_plus-L", "PP-DocLayout-L"]
VL_SAME_LINES = "Same text lines as classic (fair comparison)"
VL_LAYOUT = "VL's own layout zones (text / table / formula)"
RENDER_DPI = 150
MAX_SIDE = 2200 # cap the long side of a rendered page (px)
MAX_PAGES = 30
MAX_VL_ITEMS = 150 # bound GPU time per page
SAMPLE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples", "sample.png")
# Free ZeroGPU caps each call's duration and each visitor's GPU-seconds per day,
# so recognition runs in several short allocations rather than one long one.
GPU_DURATION = 55
CHUNK_LINES = 24 # short single-line crops
CHUNK_ZONES = 10 # whole regions (can be long)
PROMPTS = {"ocr": "OCR:", "table": "Table Recognition:",
"formula": "Formula Recognition:", "chart": "Chart Recognition:"}
MAX_TOKENS = {"ocr": 1024, "table": 1536, "formula": 512, "chart": 1024}
LABEL_TO_TASK = {"table": "table", "formula": "formula", "equation": "formula",
"chart": "chart"}
ON_ZEROGPU = bool(os.environ.get("SPACES_ZERO_GPU"))
USE_CUDA = torch is not None and (ON_ZEROGPU or torch.cuda.is_available())
DEVICE = "cuda" if USE_CUDA else "cpu"
# ---------------------------------------------------------------------------
# Models (lazy singletons)
# ---------------------------------------------------------------------------
_classic: dict = {}
_layout = {"model": None}
_vlm = {"model": None, "processor": None, "error": None}
def _get_classic(size: str):
if size in _classic:
return _classic[size]
from paddleocr import PaddleOCR
det, rec = CLASSIC_MODELS[size]
print(f"[classic] loading {det} + {rec}", flush=True)
engine = PaddleOCR(
text_detection_model_name=det,
text_recognition_model_name=rec,
# Keep the page exactly as rendered so boxes line up with the image.
use_doc_orientation_classify=False,
use_doc_unwarping=False,
use_textline_orientation=False,
)
_classic[size] = engine
return engine
def _get_layout():
if _layout["model"] is not None:
return _layout["model"]
from paddleocr import LayoutDetection
last = None
for name in LAYOUT_MODELS:
try:
_layout["model"] = LayoutDetection(model_name=name)
print("[layout] using", name, flush=True)
return _layout["model"]
except Exception as e: # noqa: BLE001
last = e
print(f"[layout] {name} failed: {e}", flush=True)
raise RuntimeError(f"Could not load a layout model: {last}")
def _patch_transformers_masking():
"""PaddleOCR-VL's remote code calls create_causal_mask(inputs_embeds=...),
but transformers 4.55 names that parameter `input_embeds`. Rename it."""
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
@functools.wraps(orig)
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
def _load_vlm():
if _vlm["model"] is not None:
return
if torch is None:
raise RuntimeError("PyTorch is not installed, so PaddleOCR-VL is unavailable here.")
_patch_transformers_masking()
from transformers import AutoModelForCausalLM, AutoProcessor
print("[vlm] loading", VL_MODEL, flush=True)
model = AutoModelForCausalLM.from_pretrained(
VL_MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16
).eval()
if DEVICE == "cuda":
model = model.to("cuda")
_vlm["processor"] = AutoProcessor.from_pretrained(VL_MODEL, trust_remote_code=True)
_vlm["model"] = model
print("[vlm] ready", flush=True)
# Load the VLM at import so ZeroGPU can pack the weights (ZeroGPU best practice).
try:
_load_vlm()
except Exception as e: # noqa: BLE001
_vlm["error"] = str(e)
print("[vlm] eager load failed, will retry on first use:", e, flush=True)
# ---------------------------------------------------------------------------
# Pages
# ---------------------------------------------------------------------------
def _cap(img: Image.Image) -> Image.Image:
w, h = img.size
if max(w, h) > MAX_SIDE:
s = MAX_SIDE / max(w, h)
img = img.resize((int(w * s), int(h * s)), Image.LANCZOS)
return img
def load_pages(path):
if not path:
return []
if os.path.splitext(path)[1].lower() == ".pdf":
doc = fitz.open(path)
zoom = RENDER_DPI / 72.0
pages = []
for i, page in enumerate(doc):
if i >= MAX_PAGES:
break
pix = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
pages.append(_cap(Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")))
doc.close()
return pages
return [_cap(Image.open(path).convert("RGB"))]
# ---------------------------------------------------------------------------
# Engine 1: classic PaddleOCR (CPU)
# ---------------------------------------------------------------------------
def _result_dict(res):
"""PaddleOCR 3.x results are dict-like; .json gives plain lists."""
try:
j = res.json
if isinstance(j, dict):
return j.get("res", j)
except Exception: # noqa: BLE001
pass
return res
def run_classic(img: Image.Image, size: str):
"""-> (items, detected_boxes). detected_boxes includes lines the recognizer
returned empty, so the VL 'same lines' mode sees every detection."""
engine = _get_classic(size)
bgr = np.ascontiguousarray(np.array(img.convert("RGB"))[:, :, ::-1]) # Paddle expects BGR
res = _result_dict(list(engine.predict(bgr))[0])
texts = list(res.get("rec_texts") or [])
scores = list(res.get("rec_scores") or [])
polys = list(res.get("rec_polys") or [])
items = []
for i, text in enumerate(texts):
if i >= len(polys) or not str(text).strip():
continue
items.append({
"text": str(text).strip(),
"box": quad_to_box(polys[i]),
"conf": float(scores[i]) * 100 if i < len(scores) else None,
"label": "text line",
})
det = [quad_to_box(p) for p in (res.get("dt_polys") or polys)]
return reading_order(items), det
# ---------------------------------------------------------------------------
# Engine 2: PaddleOCR-VL (ZeroGPU)
# ---------------------------------------------------------------------------
def _vlm_one(crop: Image.Image, task: str, max_tokens: int) -> str:
model, processor = _vlm["model"], _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 new tokens; otherwise the prompt is echoed back.
text = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:],
skip_special_tokens=True)[0]
for marker in ("Assistant:", "assistant\n", "assistant:"):
if marker in text:
text = text.split(marker)[-1]
return text.strip()
@spaces.GPU(duration=GPU_DURATION)
def _recognize_chunk(img: Image.Image, chunk: list[dict]) -> list[dict]:
"""One short GPU allocation. Returns the chunk (ZeroGPU runs this in a
separate process, so in-place mutation would not come back)."""
_load_vlm()
W, H = img.size
for z in chunk:
x1, y1, x2, y2 = z["box"]
if z.get("line"):
# a little vertical padding helps the recognizer on tight line boxes
pad = max(2, int((y2 - y1) * 0.15))
y1, y2 = max(0, y1 - pad), min(H, y2 + pad)
max_tokens = 128
else:
max_tokens = MAX_TOKENS.get(z["task"], 1024)
try:
z["text"] = _vlm_one(img.crop((x1, y1, x2, y2)), z["task"], max_tokens)
except Exception as e: # noqa: BLE001
z["text"] = f"[recognition error: {e}]"
return chunk
def _layout_zones(img: Image.Image) -> list[dict]:
out = _get_layout().predict(np.array(img), batch_size=1)
res = _result_dict(list(out)[0])
W, H = img.size
zones = []
for b in res.get("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, x2, y2 = max(0, x1), max(0, y1), min(W, x2), min(H, y2)
if x2 - x1 < 3 or y2 - y1 < 3:
continue
label = str(b.get("label", "text")).lower().strip()
zones.append({"box": [int(x1), int(y1), int(x2), int(y2)], "label": label,
"task": LABEL_TO_TASK.get(label, "ocr")})
return reading_order(zones)
def run_vl(img: Image.Image, mode: str, line_boxes: list[list[int]]):
if mode == VL_SAME_LINES:
zones = [{"box": b, "label": "text line", "task": "ocr", "line": True}
for b in line_boxes]
chunk = CHUNK_LINES
else:
zones = _layout_zones(img)
chunk = CHUNK_ZONES
zones = reading_order(zones)[:MAX_VL_ITEMS]
for i in range(0, len(zones), chunk):
zones[i : i + chunk] = _recognize_chunk(img, zones[i : i + chunk])
return [{"text": z.get("text", ""), "box": z["box"], "conf": None, "label": z["label"]}
for z in zones]
# ---------------------------------------------------------------------------
# Gradio callbacks
# ---------------------------------------------------------------------------
EMPTY_TABLE = gr.update(value=[], headers=TABLE_HEADERS)
def on_upload(path):
pages = load_pages(path)
if not pages:
return (gr.update(choices=[], value=None), [], None, placeholder_html("Upload a PDF or image."),
placeholder_html("Upload a PDF or image."), [], EMPTY_TABLE, "", "[]", None)
choices = [(f"Page {i + 1}", i) for i in range(len(pages))]
msg = f"Loaded {len(pages)} page(s). Pick a page and click **Run both**."
ready = placeholder_html("Click “Run both” to OCR this page.")
return (gr.update(choices=choices, value=0), pages, path, ready, ready, [], EMPTY_TABLE, msg,
"[]", None)
def load_sample():
return on_upload(SAMPLE_PATH if os.path.exists(SAMPLE_PATH) else None)
def _summary(t_classic, n_classic, t_vl, n_vl, stats, vl_mode, note=""):
lines = [f"**PaddleOCR (classic):** {n_classic} lines in {t_classic:.1f}s"]
if t_vl is not None:
mode = "same lines" if vl_mode == VL_SAME_LINES else "layout zones"
lines.append(f"**PaddleOCR-VL:** {n_vl} boxes ({mode}) in {t_vl:.1f}s")
if stats:
c = stats["counts"]
mean = f"{stats['mean_ratio'] * 100:.1f}%" if stats["mean_ratio"] is not None else "—"
lines.append(
f"**Agreement:** {c['same']} same · {c['minor']} minor · {c['differs']} differ · "
f"{c['only-classic']} only classic · {c['only-vl']} only VL "
f"(mean similarity {mean})"
)
if note:
lines.append(note)
return " \n".join(lines)
def run_both(page_idx, pages, classic_size, vl_mode, use_vl, only_diff):
if not pages:
raise gr.Error("Upload a document (or click “Load sample”) first.")
pno = int(page_idx or 0)
img = pages[pno]
uri = image_to_data_uri(img) # encode once, reuse for both panels
waiting = placeholder_html("Waiting for classic PaddleOCR…")
yield (placeholder_html("Running classic PaddleOCR on CPU…"), waiting, [],
EMPTY_TABLE, "Running classic PaddleOCR…")
t0 = time.perf_counter()
try:
classic, det_boxes = run_classic(img, classic_size)
except Exception as e: # noqa: BLE001
traceback.print_exc()
raise gr.Error(f"Classic PaddleOCR failed: {e}")
t_classic = time.perf_counter() - t0
classic_html = overlay_html(img, classic, "classic", uri, pno)
if not use_vl:
yield (classic_html, placeholder_html("PaddleOCR-VL is switched off."), [],
EMPTY_TABLE, _summary(t_classic, len(classic), None, 0, None, vl_mode))
return
if vl_mode == VL_SAME_LINES and not det_boxes:
yield (classic_html, placeholder_html("Classic found no text lines to re-read."),
[], EMPTY_TABLE, _summary(t_classic, len(classic), None, 0, None, vl_mode))
return
yield (classic_html, placeholder_html("Running PaddleOCR-VL on ZeroGPU…"), [],
EMPTY_TABLE,
_summary(t_classic, len(classic), None, 0, None, vl_mode) + " \nRunning PaddleOCR-VL…")
t0 = time.perf_counter()
try:
vl = run_vl(img, vl_mode, det_boxes)
except Exception as e: # noqa: BLE001
traceback.print_exc()
yield (classic_html, placeholder_html(f"PaddleOCR-VL failed: {e}"), [], EMPTY_TABLE,
_summary(t_classic, len(classic), None, 0, None, vl_mode))
return
t_vl = time.perf_counter() - t0
rows, stats = compare(classic, vl) # also tags boxes with their diff row
note = ""
if vl_mode == VL_SAME_LINES and len(det_boxes) > MAX_VL_ITEMS:
note = f"_VL re-read the first {MAX_VL_ITEMS} of {len(det_boxes)} lines (GPU budget)._"
yield (
overlay_html(img, classic, "classic", uri, pno),
overlay_html(img, vl, "vl", uri, pno),
rows,
gr.update(value=table_rows(rows, only_diff), headers=TABLE_HEADERS),
_summary(t_classic, len(classic), t_vl, len(vl), stats, vl_mode, note),
)
def on_filter(rows, only_diff):
return gr.update(value=table_rows(rows or [], only_diff), headers=TABLE_HEADERS)
def export_pdf(edits_json, src_path, pages):
if not src_path or not pages:
raise gr.Error("Upload a document (or click “Load sample”) first.")
edits = parse_edits(edits_json)
if not edits:
raise gr.Error("No saved edits yet. Click a box, change the text, then press Enter or Save.")
try:
path, applied, skipped = build_edited_pdf(src_path, pages, edits)
except Exception as e: # noqa: BLE001
traceback.print_exc()
raise gr.Error(f"Could not build the PDF: {e}")
msg = f"Edited PDF ready: {applied} change(s) applied"
if skipped:
msg += f", {skipped} skipped (text too long for its box, or on an unloaded page)"
return path, msg + "."
# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="OCR Compare", theme=gr.themes.Soft(), css=CSS, head=HEAD_JS) as demo:
pages_state = gr.State([])
rows_state = gr.State([])
src_state = gr.State(None)
gr.Markdown(
"# 🔍 OCR Compare — PaddleOCR vs PaddleOCR-VL\n"
"Upload a scan, run both engines, and compare them side by side. "
"**Every word or line is an invisible box — click one to edit its text, then "
"download the edited PDF.** The other engine's reading for the same spot is "
"shown under the editor."
)
with gr.Row():
file_in = gr.File(label="PDF or image", type="filepath",
file_types=[".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tif", ".tiff"],
scale=3)
with gr.Column(scale=2):
page_sel = gr.Dropdown(label="Page", choices=[], interactive=True)
with gr.Row():
sample_btn = gr.Button("Load sample")
run_btn = gr.Button("Run both", variant="primary")
with gr.Accordion("Options", open=False):
with gr.Row():
classic_size = gr.Radio(list(CLASSIC_MODELS), value="mobile (fast)",
label="Classic PaddleOCR models")
vl_mode = gr.Radio([VL_SAME_LINES, VL_LAYOUT], value=VL_SAME_LINES,
label="PaddleOCR-VL boxes",
info="Same lines isolates recognition quality; layout "
"zones shows how VL would read the page on its own.")
use_vl = gr.Checkbox(value=True, label="Run PaddleOCR-VL (uses ZeroGPU quota)")
with gr.Row():
show_boxes = gr.Checkbox(value=False, label="Show all boxes (colored by agreement)")
show_text = gr.Checkbox(value=False, label="Show recognized text on the page")
status = gr.Markdown("Upload a PDF or image, or click **Load sample**.")
with gr.Row(equal_height=False):
with gr.Column():
gr.Markdown("### PaddleOCR (classic)")
classic_panel = gr.HTML(placeholder_html("No page yet."), padding=False)
with gr.Column():
gr.Markdown("### PaddleOCR-VL")
vl_panel = gr.HTML(placeholder_html("No page yet."), padding=False)
with gr.Row(equal_height=True):
with gr.Column(scale=3):
gr.HTML('<span id="ocrc-count">No edits yet</span>', padding=False)
export_msg = gr.Markdown("")
export_btn = gr.Button("⬇️ Download edited PDF", variant="primary", scale=1)
pdf_out = gr.File(label="Edited PDF", scale=2)
# Filled by the page script (overlay.HEAD_JS) with the saved edits as JSON;
# hidden with CSS rather than visible=False so it stays in the DOM.
edits_box = gr.Textbox("[]", elem_id="ocrc-edits", show_label=False, container=False)
gr.Markdown(
"### Differences\n"
"Green = identical · yellow = minor · red = different · blue/purple = found by "
"one engine only. Row numbers match the “diff row” shown when you click a box."
)
only_diff = gr.Checkbox(value=True, label="Only show rows that differ")
diff_table = gr.Dataframe(headers=TABLE_HEADERS, value=[], wrap=True, interactive=False,
column_widths=["6%", "14%", "8%", "36%", "36%"])
# --- wiring ---
upload_out = [page_sel, pages_state, src_state, classic_panel, vl_panel, rows_state,
diff_table, status, edits_box, pdf_out]
clear_edits_js = "() => { window.__ocrcClear && window.__ocrcClear(); }"
file_in.change(on_upload, inputs=file_in, outputs=upload_out)
sample_btn.click(load_sample, inputs=None, outputs=upload_out)
# A new document starts with no edits (browser-side store).
file_in.change(None, js=clear_edits_js)
sample_btn.click(None, js=clear_edits_js)
export_btn.click(export_pdf, inputs=[edits_box, src_state, pages_state],
outputs=[pdf_out, export_msg], concurrency_limit=4)
def _page_changed(_):
ready = placeholder_html("Click “Run both” to OCR this page.")
return ready, ready, [], EMPTY_TABLE
page_sel.change(_page_changed, inputs=page_sel,
outputs=[classic_panel, vl_panel, rows_state, diff_table])
run_btn.click(
run_both,
inputs=[page_sel, pages_state, classic_size, vl_mode, use_vl, only_diff],
outputs=[classic_panel, vl_panel, rows_state, diff_table, status],
)
only_diff.change(on_filter, inputs=[rows_state, only_diff], outputs=diff_table)
# Display toggles are pure CSS classes on <body>; no server round trip.
show_boxes.change(None, inputs=show_boxes, outputs=None,
js="(v) => { document.body.classList.toggle('ocrc-show-boxes', v); }")
show_text.change(None, inputs=show_text, outputs=None,
js="(v) => { document.body.classList.toggle('ocrc-show-text', v); }")
demo.queue(default_concurrency_limit=1, max_size=16)
if __name__ == "__main__":
demo.launch()