| """grinpress v3 OCR evaluation. |
| |
| python eval_ocr.py refs # Phase 2: tesseract v2 / tesseract v3 / CRNN on the 48 reference renders |
| python eval_ocr.py models --unet onnx/unet_distilled_int8.onnx --text onnx/text_encoder_int8.onnx [--teacher ckpt] |
| # Phase 5: 48-grid at 2 and 4 steps (+14-grid teacher vs student) -> assets/eval_ocr_v3.json, eval_compare_v3.json, |
| # samples_4step_v3.png, samples_2step_v3.png, teacher_vs_student_v3.png |
| |
| Scores (per string): exact = case-sensitive exact match after whitespace normalisation; exact_ci = v2's metric |
| (upper-cased, stray face-edge glyph stripped); char_acc = 1 - Levenshtein/len on the exact_ci-normalised strings |
| (v2 definition), char_acc_cs = same on case-sensitive strings. Primary v3 metric: CRNN exact (case-sensitive). |
| """ |
| import argparse, json, os, random, re, sys, time |
| import numpy as np, torch |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, ROOT) |
|
|
|
|
| def lev(a, b): |
| d = list(range(len(b) + 1)) |
| for i, ca in enumerate(a, 1): |
| p, d[0] = d[0], i |
| for j, cb in enumerate(b, 1): |
| p, d[j] = d[j], min(d[j] + 1, d[j - 1] + 1, p + (ca != cb)) |
| return d[len(b)] |
|
|
|
|
| def norm_ci(s): |
| s = re.sub(r"\s+", " ", s.strip().upper()) |
| return re.sub(r"^[)|}\].,'\u2018\u201c-] ", "", s) |
|
|
|
|
| def norm_cs(s): |
| return re.sub(r"\s+", " ", s.strip()) |
|
|
|
|
| def score_text(pred, gt): |
| p, g = norm_ci(pred), norm_ci(gt) |
| pc, gc = norm_cs(pred), norm_cs(gt) |
| return dict(exact=pc == gc, exact_ci=p == g, char_acc=max(0.0, 1 - lev(p, g) / max(1, len(g))), |
| char_acc_cs=max(0.0, 1 - lev(pc, gc) / max(1, len(gc)))) |
|
|
|
|
| def summarize(rows, key): |
| P = [r[key] for r in rows] |
| return dict(n=len(P), exact=sum(p["exact"] for p in P) / len(P), exact_ci=sum(p["exact_ci"] for p in P) / len(P), |
| char_acc=float(np.mean([p["char_acc"] for p in P])), char_acc_cs=float(np.mean([p["char_acc_cs"] for p in P]))) |
|
|
|
|
| def fonts(): |
| from generate_data import ensure_fonts |
| return sorted(ensure_fonts()) |
|
|
|
|
| def reference_images(grid): |
| from generate_data_v2 import render_v2 |
| F = fonts() |
| return [render_v2(128, random.Random(seed), F, **{k: sp[k] for k in ("color", "expr", "text", "layout", "paper")})[0] |
| for sp, _, seed in grid] |
|
|
|
|
| def ocr_all(imgs, layouts, crnn=None, tess=("v3",)): |
| from crnn import recognize |
| from ocr_tess import tess_v2, tess_v3 |
| out = {} |
| if crnn is not None: |
| out["crnn"] = recognize(crnn, imgs, layouts) |
| if "v2" in tess: |
| out["tess_v2"] = [tess_v2(im, l) for im, l in zip(imgs, layouts)] |
| if "v3" in tess: |
| out["tess_v3"] = [tess_v3(im, l) for im, l in zip(imgs, layouts)] |
| if "v3s" in tess: |
| out["tess_v3_sauvola"] = [tess_v3(im, l, "sauvola") for im, l in zip(imgs, layouts)] |
| return out |
|
|
|
|
| FONT = ImageFont.load_default(size=13) |
|
|
|
|
| def grid_png(imgs, labels, cols, path, cell=192, title=None): |
| rows = (len(imgs) + cols - 1) // cols; lh = 34; th = 30 if title else 0 |
| G = Image.new("RGB", (cols * cell, th + rows * (cell + lh)), "white"); d = ImageDraw.Draw(G) |
| if title: d.text((6, 6), title, fill="black", font=ImageFont.load_default(size=16)) |
| for k, (im, lab) in enumerate(zip(imgs, labels)): |
| r, c = divmod(k, cols); x0, y0 = c * cell, th + r * (cell + lh) |
| G.paste(im.resize((cell, cell), Image.NEAREST), (x0, y0)) |
| for li, line in enumerate(lab.split("\n")[:2]): |
| d.text((x0 + 3, y0 + cell + 2 + li * 15), line, fill="black", font=FONT) |
| G.save(path); return path |
|
|
|
|
| |
| def cmd_refs(a): |
| from eval_grid import grid48 |
| from crnn import load_crnn |
| g = grid48(); imgs = reference_images(g); lays = [sp["layout"] for sp, _, _ in g]; texts = [sp["text"] for sp, _, _ in g] |
| crnn = load_crnn(a.crnn) if a.crnn and os.path.exists(a.crnn) else None |
| t0 = time.time() |
| preds = ocr_all(imgs, lays, crnn, tess=("v2", "v3", "v3s")) |
| res = {"what": "OCR on the 48 procedural reference renders (generate_data_v2.render_v2, seeds 1000+i, byte-identical to v2)", |
| "crnn": a.crnn if crnn is not None else None, "per_sample": [], "summary": {}} |
| for i, (sp, cap, seed) in enumerate(g): |
| row = dict(text=sp["text"], layout=sp["layout"], paper=sp["paper"], expr=sp["expr"], seed=seed) |
| for k, v in preds.items(): row[k] = dict(ocr=v[i], **score_text(v[i], sp["text"])) |
| res["per_sample"].append(row) |
| for k in preds: res["summary"][k] = summarize(res["per_sample"], k) |
| res["seconds"] = round(time.time() - t0, 1) |
| os.makedirs(os.path.dirname(a.out), exist_ok=True) |
| json.dump(res, open(a.out, "w"), indent=1) |
| print(json.dumps(res["summary"], indent=1)) |
| for k in preds: |
| miss = [(r["text"], r[k]["ocr"]) for r in res["per_sample"] if not r[k]["exact"]] |
| print(k, "misses (case-sensitive):", miss[:60]) |
|
|
|
|
| |
| def ort_sess(p, threads=os.cpu_count()): |
| import onnxruntime as ort |
| so = ort.SessionOptions(); so.intra_op_num_threads = threads; so.inter_op_num_threads = 1 |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| return ort.InferenceSession(p, so, providers=["CPUExecutionProvider"]) |
|
|
|
|
| def noise(seed): |
| return torch.randn(1, 3, 128, 128, generator=torch.Generator().manual_seed(seed)).numpy() |
|
|
|
|
| def sched(n): |
| return [1.0 - i / n for i in range(n + 1)] |
|
|
|
|
| def student_onnx(te, un, caps, seeds, steps, bs=8): |
| from conditioning import tokenize |
| out = [] |
| for b in range(0, len(caps), bs): |
| c, sd = caps[b:b + bs], seeds[b:b + bs] |
| ids, mk = tokenize(c) |
| ctx, cm = te.run(None, {"input_ids": ids.numpy(), "attention_mask": mk.long().numpy()}) |
| x = np.concatenate([noise(s) for s in sd]); ts = sched(steps) |
| for i in range(steps): |
| v = un.run(None, {"x_t": x, "t": np.full((len(c),), ts[i], np.float32), "context": ctx, "context_mask": cm})[0] |
| x = x - (ts[i] - ts[i + 1]) * v |
| out.append(np.clip(x, -1, 1)) |
| return np.concatenate(out) |
|
|
|
|
| @torch.no_grad() |
| def teacher_images(teacher_path, caps, seeds, steps=50, cfg=5.0, bs=8, device="cpu"): |
| from train_teacher_v3 import load_flow, sample_cfg |
| T = load_flow(teacher_path, device) |
| out = [] |
| for b in range(0, len(caps), bs): |
| z = torch.from_numpy(np.concatenate([noise(s) for s in seeds[b:b + bs]])).to(device) |
| out.append(sample_cfg(T, caps[b:b + bs], z, steps, cfg).cpu().numpy()) |
| return np.concatenate(out) |
|
|
|
|
| def to_img(x): |
| return Image.fromarray(((x.transpose(1, 2, 0) + 1) * 127.5).round().clip(0, 255).astype(np.uint8)) |
|
|
|
|
| def cmd_models(a): |
| from eval_grid import grid48, grid14 |
| from crnn import load_crnn |
| crnn = load_crnn(a.crnn) |
| tess = tuple(t for t in a.tess.split(",") if t) |
| os.makedirs(a.assets, exist_ok=True) |
| te, un = ort_sess(a.text), ort_sess(a.unet) |
| g = grid48(); caps = [c for _, c, _ in g]; seeds = [s for *_, s in g]; lays = [sp["layout"] for sp, _, _ in g] |
| imgs, res = {}, {"unet": os.path.basename(a.unet), "text_encoder": os.path.basename(a.text), "n": len(g), |
| "captions_truncated": 0, "primary_metric": "crnn exact (case-sensitive)", "per_sample": [], "summary": {}} |
| from conditioning import MAX_LEN |
| res["captions_truncated"] = sum(len(c) > MAX_LEN for c in caps) |
| t0 = time.time() |
| for steps in a.steps: |
| imgs[f"s{steps}"] = [to_img(z) for z in student_onnx(te, un, caps, seeds, steps)] |
| imgs["ref"] = reference_images(g) |
| preds = {k: ocr_all(v, lays, crnn, tess) for k, v in imgs.items()} |
| for i, (sp, cap, seed) in enumerate(g): |
| row = dict(sp, caption=cap, seed=seed) |
| for k in imgs: |
| row[k] = {m: dict(ocr=p[i], **score_text(p[i], sp["text"])) for m, p in preds[k].items()} |
| res["per_sample"].append(row) |
| for k in imgs: |
| res["summary"][k] = {m: summarize([{"x": r[k][m]} for r in res["per_sample"]], "x") for m in preds[k]} |
| res["seconds"] = round(time.time() - t0, 1) |
| json.dump(res, open(os.path.join(a.assets, "eval_ocr_v3.json"), "w"), indent=1) |
| for steps in a.steps: |
| lab = [f"'{sp['text']}' -> '{r[f's{steps}']['crnn']['ocr']}'\n{sp['expr']}/{sp['layout']}/{sp['paper']}" |
| for (sp, _, _), r in zip(g, res["per_sample"])] |
| grid_png(imgs[f"s{steps}"], lab, 8, os.path.join(a.assets, f"samples_{steps}step_v3.png"), |
| title=f"grinpress v3 student ({os.path.basename(a.unet)}), {steps} steps, seeds 1000+i (label: text -> CRNN read)") |
| print(json.dumps(res["summary"], indent=1)) |
| |
| g14 = grid14(); c14 = [c for _, c, _ in g14]; s14 = [s for *_, s in g14]; l14 = [sp["layout"] for sp, _, _ in g14] |
| rows = {"ref": reference_images(g14)} |
| if a.teacher: |
| rows["teacher"] = [to_img(z) for z in teacher_images(a.teacher, c14, s14, a.teacher_steps, a.cfg, device=a.device)] |
| for steps in a.steps: |
| rows[f"s{steps}"] = [to_img(z) for z in student_onnx(te, un, c14, s14, steps)] |
| cmp = {} |
| for k, ims in rows.items(): |
| p = ocr_all(ims, l14, crnn, tess) |
| cmp[k] = {m: dict(summarize([{"x": score_text(o, sp["text"])} for o, (sp, _, _) in zip(v, g14)], "x"), ocr=v) for m, v in p.items()} |
| json.dump(cmp, open(os.path.join(a.assets, "eval_compare_v3.json"), "w"), indent=1) |
| all_imgs, labels = [], [] |
| for k, ims in rows.items(): |
| for i, im in enumerate(ims): |
| all_imgs.append(im); labels.append(f"'{g14[i][0]['text']}' {k}\n-> '{cmp[k]['crnn']['ocr'][i]}'") |
| grid_png(all_imgs, labels, len(g14), os.path.join(a.assets, "teacher_vs_student_v3.png"), cell=160, |
| title="rows: " + " | ".join(rows) + " (same noise per column; label: CRNN read)") |
| print(json.dumps({k: {m: {x: y for x, y in v.items() if x != "ocr"} for m, v in d.items()} for k, d in cmp.items()}, indent=1)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| r = sub.add_parser("refs") |
| r.add_argument("--crnn", default=os.path.join(ROOT, "artifacts", "crnn.safetensors")) |
| r.add_argument("--out", default=os.path.join(ROOT, "assets", "eval_ocr_refs_phase2.json")) |
| m = sub.add_parser("models") |
| m.add_argument("--unet", required=True) |
| m.add_argument("--text", required=True) |
| m.add_argument("--teacher", default="") |
| m.add_argument("--teacher_steps", type=int, default=50) |
| m.add_argument("--cfg", type=float, default=5.0) |
| m.add_argument("--steps", type=int, nargs="+", default=[4, 2]) |
| m.add_argument("--crnn", default=os.path.join(ROOT, "artifacts", "crnn.safetensors")) |
| m.add_argument("--tess", default="v3", help="comma list of tesseract variants to also run (v2,v3,v3s) or ''") |
| m.add_argument("--assets", default=os.path.join(ROOT, "assets")) |
| m.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| a = ap.parse_args() |
| torch.set_num_threads(os.cpu_count()) |
| {"refs": cmd_refs, "models": cmd_models}[a.cmd](a) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
