242 lines
9.7 KiB
Python
242 lines
9.7 KiB
Python
#!/usr/bin/env python3
|
||
"""Benchmark OCR engines on pages of a native-text PDF, using its text layer as ground truth.
|
||
|
||
python3 scripts/bench_ocr.py --pdf source.pdf --pages 12,20,30 --engine tesseract \
|
||
--engine "vision:http://localhost:8098|baidu/Unlimited-OCR" --engine docling --engine paddleocr
|
||
|
||
The PDF's own text layer is what a perfect transcription of the rendered page would say, so no
|
||
hand-labelled data is needed. Two scores per page, both on lower-cased word tokens:
|
||
* bag F1 -- order-insensitive overlap (robust to differing column reading order)
|
||
* seq sim -- difflib ratio over the token sequence (penalises scrambled order)
|
||
Also wall-clock seconds per page and the peak resident memory of THIS process (a remote engine's own
|
||
footprint is not visible here; read it from `docker stats`).
|
||
|
||
Engines are optional: one whose library or server is missing is reported as unavailable, not an error.
|
||
Read-only: renders into a temp directory and writes nothing else unless --out is given.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import difflib
|
||
import json
|
||
import re
|
||
import resource
|
||
import statistics
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
|
||
_TOKEN = re.compile(r"[\wÄÖÜäöüß]+", re.UNICODE)
|
||
|
||
|
||
def tokens(text: str) -> list[str]:
|
||
# Hyphenation at line ends ("zu-\nführen") is joined so it does not count against an engine
|
||
# that correctly re-joins it (or for one that keeps the break).
|
||
text = re.sub(r"-\s*\n\s*", "", text or "")
|
||
return [t.lower() for t in _TOKEN.findall(text)]
|
||
|
||
|
||
def bag_f1(ref: list[str], hyp: list[str]) -> float:
|
||
if not ref or not hyp:
|
||
return 0.0
|
||
overlap = sum((Counter(ref) & Counter(hyp)).values())
|
||
if not overlap:
|
||
return 0.0
|
||
p, r = overlap / len(hyp), overlap / len(ref)
|
||
return 2 * p * r / (p + r)
|
||
|
||
|
||
def seq_sim(ref: list[str], hyp: list[str]) -> float:
|
||
if not ref or not hyp:
|
||
return 0.0
|
||
return difflib.SequenceMatcher(None, ref, hyp, autojunk=False).ratio()
|
||
|
||
|
||
def render(pdf: Path, page: int, dpi: int, out: Path) -> Path:
|
||
import pymupdf
|
||
with pymupdf.open(str(pdf)) as doc:
|
||
doc[page - 1].get_pixmap(dpi=dpi).save(str(out))
|
||
return out
|
||
|
||
|
||
def truth(pdf: Path, page: int) -> str:
|
||
import pymupdf
|
||
with pymupdf.open(str(pdf)) as doc:
|
||
return doc[page - 1].get_text()
|
||
|
||
|
||
def pick_pages(pdf: Path, count: int, min_words: int) -> list[int]:
|
||
"""Evenly spaced pages with enough text to score meaningfully (adverts and photo pages drop out)."""
|
||
import pymupdf
|
||
with pymupdf.open(str(pdf)) as doc:
|
||
eligible = [n + 1 for n in range(len(doc)) if len(tokens(doc[n].get_text())) >= min_words]
|
||
if len(eligible) <= count:
|
||
return eligible
|
||
step = len(eligible) / count
|
||
return [eligible[int(i * step)] for i in range(count)]
|
||
|
||
|
||
# ── engines: each is a callable (png_path) -> transcribed text, or raises Unavailable ─────────────
|
||
|
||
class Unavailable(Exception):
|
||
pass
|
||
|
||
|
||
def engine_tesseract(lang: str = "deu+eng"):
|
||
import shutil
|
||
import subprocess
|
||
if not shutil.which("tesseract"):
|
||
raise Unavailable("tesseract is not installed here")
|
||
|
||
def run(png: Path) -> str:
|
||
return subprocess.run(["tesseract", str(png), "stdout", "-l", lang], capture_output=True,
|
||
text=True, timeout=180, check=True).stdout
|
||
return run
|
||
|
||
|
||
def engine_vision(spec: str):
|
||
"""`vision:<base_url>|<model>` -- an OpenAI-compatible vision server, called exactly as the app calls it."""
|
||
try:
|
||
base, model = spec.split("|", 1)
|
||
except ValueError:
|
||
raise Unavailable("use vision:<url>|<model>")
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
try:
|
||
from core import magazine as mag
|
||
except Exception as e: # the app's own module needs its environment
|
||
raise Unavailable(f"cannot import the app's OCR client ({type(e).__name__})")
|
||
import requests
|
||
try:
|
||
requests.get(base.rstrip("/") + ("/v1/models" if "/v1" not in base else "/models"), timeout=5).raise_for_status()
|
||
except Exception as e:
|
||
raise Unavailable(f"server not reachable ({type(e).__name__})")
|
||
|
||
def run(png: Path) -> str:
|
||
regions = mag.vision_ocr_page(png, base, model, timeout=300.0)
|
||
if regions is None:
|
||
raise RuntimeError("engine returned nothing")
|
||
return "\n".join(str(r.get("text") or "") for r in regions)
|
||
return run
|
||
|
||
|
||
def engine_docling():
|
||
try:
|
||
from docling.document_converter import DocumentConverter
|
||
except Exception as e:
|
||
raise Unavailable(f"docling not installed ({type(e).__name__})")
|
||
conv = DocumentConverter()
|
||
|
||
def run(png: Path) -> str:
|
||
return conv.convert(str(png)).document.export_to_text()
|
||
return run
|
||
|
||
|
||
def engine_paddleocr(lang: str = "de"):
|
||
try:
|
||
from paddleocr import PaddleOCR
|
||
except Exception as e:
|
||
raise Unavailable(f"paddleocr not installed ({type(e).__name__})")
|
||
try: # 3.x: page orientation/unwarping models are extra work a straight scan does not need
|
||
ocr = PaddleOCR(lang=lang, use_doc_orientation_classify=False, use_doc_unwarping=False,
|
||
use_textline_orientation=False)
|
||
except TypeError: # 2.x
|
||
ocr = PaddleOCR(lang=lang, use_angle_cls=False)
|
||
|
||
def run(png: Path) -> str:
|
||
lines: list[str] = []
|
||
res = ocr.predict(str(png)) if hasattr(ocr, "predict") else ocr.ocr(str(png))
|
||
for page in res or []:
|
||
if isinstance(page, dict) or hasattr(page, "get"): # 3.x result object
|
||
lines.extend(str(t) for t in (page.get("rec_texts") or []))
|
||
continue
|
||
for item in page or []: # 2.x [[box, (text, score)], ...]
|
||
try:
|
||
lines.append(item[1][0])
|
||
except (IndexError, TypeError):
|
||
pass
|
||
return "\n".join(lines)
|
||
return run
|
||
|
||
|
||
def build_engine(spec: str):
|
||
if spec == "tesseract":
|
||
return engine_tesseract()
|
||
if spec.startswith("vision:"):
|
||
return engine_vision(spec[len("vision:"):])
|
||
if spec == "docling":
|
||
return engine_docling()
|
||
if spec == "paddleocr":
|
||
return engine_paddleocr()
|
||
raise Unavailable(f"unknown engine {spec!r}")
|
||
|
||
|
||
def peak_rss_mb() -> float:
|
||
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
|
||
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
ap.add_argument("--pdf", required=True, type=Path)
|
||
ap.add_argument("--pages", default="", help="comma-separated 1-based page numbers (default: evenly spaced)")
|
||
ap.add_argument("--count", type=int, default=8, help="pages to sample when --pages is not given")
|
||
ap.add_argument("--min-words", type=int, default=250)
|
||
ap.add_argument("--dpi", type=int, default=110, help="110 matches the app's ingest default")
|
||
ap.add_argument("--engine", action="append", required=True)
|
||
ap.add_argument("--out", type=Path, help="write the full results as JSON")
|
||
args = ap.parse_args()
|
||
|
||
pages = [int(p) for p in args.pages.split(",") if p.strip()] or pick_pages(args.pdf, args.count, args.min_words)
|
||
print(f"pages: {pages}")
|
||
results = []
|
||
with tempfile.TemporaryDirectory(prefix="bench-ocr-") as tmp:
|
||
pngs = {p: render(args.pdf, p, args.dpi, Path(tmp) / f"p{p}.png") for p in pages}
|
||
refs = {p: tokens(truth(args.pdf, p)) for p in pages}
|
||
for spec in args.engine:
|
||
entry = {"engine": spec, "pages": []}
|
||
try:
|
||
base_rss = peak_rss_mb()
|
||
t0 = time.perf_counter()
|
||
run = build_engine(spec)
|
||
entry["load_sec"] = round(time.perf_counter() - t0, 1)
|
||
except Unavailable as e:
|
||
entry["unavailable"] = str(e)
|
||
results.append(entry)
|
||
print(f"{spec}: unavailable — {e}")
|
||
continue
|
||
for p in pages:
|
||
t0 = time.perf_counter()
|
||
try:
|
||
hyp = tokens(run(pngs[p]))
|
||
err = None
|
||
except Exception as e:
|
||
hyp, err = [], f"{type(e).__name__}: {e}"[:200]
|
||
sec = time.perf_counter() - t0
|
||
entry["pages"].append({
|
||
"page": p, "sec": round(sec, 1), "ref_words": len(refs[p]), "hyp_words": len(hyp),
|
||
"f1": round(bag_f1(refs[p], hyp), 4), "seq": round(seq_sim(refs[p], hyp), 4), "error": err})
|
||
print(f" {spec} p{p}: f1={entry['pages'][-1]['f1']:.3f} seq={entry['pages'][-1]['seq']:.3f} {sec:.1f}s"
|
||
+ (f" ERROR {err}" if err else ""))
|
||
entry["process_rss_growth_mb"] = round(peak_rss_mb() - base_rss)
|
||
results.append(entry)
|
||
|
||
print("\n| engine | pages ok | mean F1 | mean seq | median s/page | RSS growth MB |")
|
||
print("|---|---:|---:|---:|---:|---:|")
|
||
for e in results:
|
||
ok = [r for r in e.get("pages", []) if not r["error"]]
|
||
if not ok:
|
||
print(f"| {e['engine']} | 0/{len(e.get('pages', []))} | – | – | – | – | {e.get('unavailable', '')}")
|
||
continue
|
||
print(f"| {e['engine']} | {len(ok)}/{len(e['pages'])} | {statistics.mean(r['f1'] for r in ok):.3f} | "
|
||
f"{statistics.mean(r['seq'] for r in ok):.3f} | {statistics.median(r['sec'] for r in ok):.1f} | "
|
||
f"{e.get('process_rss_growth_mb', 0)} |")
|
||
if args.out:
|
||
args.out.write_text(json.dumps({"pdf": args.pdf.name, "pages": pages, "dpi": args.dpi, "results": results},
|
||
ensure_ascii=False, indent=1), "utf-8")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|