| """ |
| benchmark.py - word-level OCR benchmark on handwritten word crops (used by notebook CELL 6). |
| |
| Recreated from how the notebook calls it; the original lived only in Colab's /content. |
| |
| Data layout (bench_sample.zip unpacks to this): |
| bench_sample/<lang>/<image> lang = hi-IN | kn-IN | ml-IN | ta-IN | te-IN |
| bench_sample/<lang>/labels.csv file,text (also .tsv / .txt with "file<TAB>text") |
| |
| python benchmark.py sample <dataset_root> --per-lang 500 # -> bench_sample.zip (same layout, random subset) |
| python benchmark.py summary bench_direct.csv bench_direct_langhint_remap.csv |
| |
| Document benchmark (full pages / PDFs against a ground-truth .txt with the same name): |
| bench_docs/<doc>.<jpg|png|pdf|...> + bench_docs/<doc>.txt (UTF-8, the whole expected text) |
| bench_docs/<lang>/<doc>.* + .txt optional: a language folder sets the language for script checks |
| |
| In the notebook: |
| run(predict_fn, "/content/bench_sample", tag="direct", limit=None) # -> bench_direct.csv + table |
| summary(["bench_direct.csv", "bench_direct_remap.csv"]) |
| upload_docs("/content/bench_docs") # Colab file picker: docs + .txt |
| run_docs(predict_doc, "/content/bench_docs", tag="docs", lang="kn-IN") # -> bench_docs.csv + table |
| """ |
| from __future__ import annotations |
|
|
| import csv, io, random, re, sys, time, unicodedata, zipfile |
| from pathlib import Path |
|
|
| from validate import cer, nfc, script_purity, wer |
|
|
| IMAGE_EXT = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp"} |
| DOC_EXT = IMAGE_EXT | {".pdf"} |
| LANGS = {"hi-IN", "kn-IN", "ml-IN", "ta-IN", "te-IN"} |
| FIELDS = ["tag", "lang", "file", "gt", "pred", "cer", "cer_norm", "word_acc", "wa_norm", "wer", |
| "empty", "halluc", "wrong_script", "seconds"] |
| ZERO_WIDTH = dict.fromkeys(map(ord, "")) |
|
|
|
|
| |
| def normalize(s: str) -> str: |
| """Lenient form for CERnorm / WAnorm: NFC, no zero-width joiners, no punctuation or spaces, |
| Indic digits -> ASCII. Separates reading errors from formatting differences.""" |
| s = nfc(s).translate(ZERO_WIDTH) |
| s = "".join(str(unicodedata.digit(c)) if unicodedata.category(c) == "Nd" else c for c in s) |
| return "".join(c for c in s if unicodedata.category(c)[0] not in ("P", "Z", "S", "C")) |
|
|
|
|
| def is_hallucination(gt: str, pred: str) -> bool: |
| """Output that is not a reading of the word: a sentence about the image, LaTeX, or far too long.""" |
| letters = [c for c in pred if unicodedata.category(c)[0] == "L"] |
| latin = sum("A" <= c <= "z" for c in letters) |
| gt_latin = any("A" <= c <= "z" for c in gt) |
| return bool(pred.strip()) and ( |
| (not gt_latin and letters and latin / len(letters) > 0.5) |
| or "$$" in pred or "\\text" in pred |
| or len(pred) > 3 * len(gt) + 10) |
|
|
|
|
| def score(gt: str, pred: str, lang: str) -> dict: |
| gt, pred = nfc(gt).strip(), nfc(pred or "").strip() |
| one = lambda s: re.sub(r"\s+", " ", s) |
| cap = lambda x: round(min(x, 1.0), 4) |
| return { |
| "cer": cap(cer(pred, gt)), |
| "cer_norm": cap(cer(normalize(pred), normalize(gt))), |
| "word_acc": int(one(pred) == one(gt)), |
| "wa_norm": int(normalize(pred) == normalize(gt)), |
| "wer": cap(wer(pred, gt)), |
| "empty": int(not pred), |
| "halluc": int(is_hallucination(gt, pred)), |
| "wrong_script": int(bool(pred) and script_purity(pred, lang) < 0.5), |
| } |
|
|
|
|
| |
| def load_labels(lang_dir: Path) -> dict[str, str]: |
| labels = {} |
| for f in sorted(lang_dir.glob("*.csv")) + sorted(lang_dir.glob("*.tsv")) + sorted(lang_dir.glob("*.txt")): |
| delim = "," if f.suffix == ".csv" else "\t" |
| for row in csv.reader(open(f, encoding="utf-8"), delimiter=delim): |
| if len(row) >= 2 and row[0] not in ("file", "filename", "image"): |
| labels[Path(row[0]).name] = row[1] |
| return labels |
|
|
|
|
| def items(root, limit=None): |
| """(lang, image path, ground truth) for every labelled image, `limit` per language.""" |
| out = [] |
| for lang_dir in sorted(p for p in Path(root).iterdir() if p.is_dir()): |
| labels = load_labels(lang_dir) |
| imgs = sorted(p for p in lang_dir.iterdir() if p.suffix.lower() in IMAGE_EXT and p.name in labels) |
| out += [(lang_dir.name, p, labels[p.name]) for p in imgs[:limit]] |
| return out |
|
|
|
|
| def doc_items(root, lang=None): |
| """(lang, document path, ground truth) for every document with a <stem>.txt next to it. |
| Language: the parent folder if it is a language code, else `lang`, else "-" (script check skipped).""" |
| out, missing = [], [] |
| for p in sorted(Path(root).rglob("*")): |
| if p.suffix.lower() not in DOC_EXT or p.name.startswith("."): |
| continue |
| gt = p.with_suffix(".txt") |
| if not gt.is_file(): |
| missing.append(p.name) |
| continue |
| out.append((p.parent.name if p.parent.name in LANGS else lang or "-", p, |
| gt.read_text(encoding="utf-8-sig"))) |
| if missing: |
| print(f"no ground truth .txt for {len(missing)} file(s), skipped: {', '.join(missing)}") |
| return out |
|
|
|
|
| def upload_docs(dest="/content/bench_docs"): |
| """Colab file picker: select documents and their .txt ground truth together (same file names).""" |
| from google.colab import files as colab_files |
| dest = Path(dest) |
| dest.mkdir(parents=True, exist_ok=True) |
| for name, data in colab_files.upload().items(): |
| if name.lower().endswith(".zip"): |
| zipfile.ZipFile(io.BytesIO(data)).extractall(dest) |
| else: |
| (dest / Path(name).name).write_bytes(data) |
| for lang, p, gt in doc_items(dest): |
| print(f"{p.relative_to(dest)} ({lang}, ground truth {len(gt)} chars)") |
| return dest |
|
|
|
|
| |
| def run(predict, root, tag="direct", limit=None, out=None): |
| """predict(path) -> text. Writes bench_<tag>.csv and prints the summary table.""" |
| return _bench(predict, items(root, limit), tag, out) |
|
|
|
|
| def run_docs(predict, root, tag="docs", lang=None, out=None): |
| """Same as run(), on whole documents: predict(path) -> the full text of the image or PDF.""" |
| return _bench(predict, doc_items(root, lang), tag, out) |
|
|
|
|
| def _bench(predict, todo, tag, out): |
| out = Path(out or f"bench_{tag}.csv") |
| with open(out, "w", encoding="utf-8", newline="") as fh: |
| wr = csv.DictWriter(fh, fieldnames=FIELDS) |
| wr.writeheader() |
| for i, (lang, path, gt) in enumerate(todo, 1): |
| t0 = time.time() |
| try: |
| pred = predict(str(path)) or "" |
| except Exception as e: |
| print(f"{path.name}: {e}") |
| pred = "" |
| row = {"tag": tag, "lang": lang, "file": path.name, "gt": gt, "pred": pred, |
| "seconds": round(time.time() - t0, 3)} |
| wr.writerow({**row, **score(gt, pred, lang)}) |
| if i % 100 == 0 or i == len(todo): |
| print(f"{i}/{len(todo)}") |
| fh.flush() |
| summary([out]) |
| return out |
|
|
|
|
| def summary(paths): |
| rows = [r for p in paths for r in csv.DictReader(open(p, encoding="utf-8"))] |
| cols = [("CER", "cer"), ("CERnorm", "cer_norm"), ("WordAcc", "word_acc"), ("WAnorm", "wa_norm"), |
| ("WER", "wer"), ("empty", "empty"), ("halluc", "halluc"), ("wrongScr", "wrong_script")] |
| print(f"{'tag':<22}{'lang':<8}{'n':>6} " + " ".join(f"{h:>8}" for h, _ in cols) + f" {'s/img':>8} {'p95 s':>8}") |
| for tag in sorted({r["tag"] for r in rows}): |
| tr = [r for r in rows if r["tag"] == tag] |
| for lang in ["ALL"] + sorted({r["lang"] for r in tr}): |
| g = tr if lang == "ALL" else [r for r in tr if r["lang"] == lang] |
| mean = lambda k: sum(float(r[k]) for r in g) / len(g) |
| secs = sorted(float(r["seconds"]) for r in g) |
| p95 = secs[min(len(secs) - 1, int(0.95 * len(secs)))] |
| print(f"{tag:<22}{lang:<8}{len(g):>6} " + " ".join(f"{mean(k):>8.3f}" for _, k in cols) |
| + f" {mean('seconds'):>8.3f} {p95:>8.3f}") |
|
|
|
|
| def sample(dataset_root, per_lang=500, out="bench_sample.zip", seed=0): |
| """Random `per_lang` crops per language from a dataset in the same layout -> zip for Colab.""" |
| rnd = random.Random(seed) |
| with zipfile.ZipFile(out, "w", zipfile.ZIP_DEFLATED) as z: |
| for lang_dir in sorted(p for p in Path(dataset_root).iterdir() if p.is_dir()): |
| labels = load_labels(lang_dir) |
| imgs = sorted(p for p in lang_dir.iterdir() if p.suffix.lower() in IMAGE_EXT and p.name in labels) |
| pick = rnd.sample(imgs, min(per_lang, len(imgs))) |
| for p in pick: |
| z.write(p, f"bench_sample/{lang_dir.name}/{p.name}") |
| buf = io.StringIO() |
| csv.writer(buf).writerows([("file", "text")] + [(p.name, labels[p.name]) for p in pick]) |
| z.writestr(f"bench_sample/{lang_dir.name}/labels.csv", buf.getvalue()) |
| print(f"{lang_dir.name}: {len(pick)}") |
| print("wrote", out) |
|
|
|
|
| if __name__ == "__main__": |
| cmd, *rest = sys.argv[1:] or ["help"] |
| if cmd == "summary": |
| summary(rest) |
| elif cmd == "sample": |
| n = int(rest[rest.index("--per-lang") + 1]) if "--per-lang" in rest else 500 |
| sample(rest[0], n) |
| else: |
| print(__doc__) |
|
|
