| """Pre-render synthetic pages to disk (clean renders; augmentation still happens at training time). |
| |
| Rendering pages is the most CPU-expensive part of the detector/orientation data pipeline. Doing it once |
| lets training run on machines with few CPU cores (e.g. Kaggle's 4-core GPU sessions). |
| |
| python -m tinydoc.data.prerender --pages 10000 --out data/synth |
| """ |
|
|
| import argparse |
| import json |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
| from ..synth.documents import KIND_WEIGHTS, build, render |
| from ..synth.fonts import FontBank |
|
|
| _bank = None |
|
|
|
|
| def _init(font_root): |
| global _bank |
| _bank = FontBank(font_root, "train") |
|
|
|
|
| def _one(args): |
| i, seed, out = args |
| rng = np.random.default_rng([seed, i]) |
| kind = str(rng.choice(list(KIND_WEIGHTS), p=list(KIND_WEIGHTS.values()))) |
| L = build(rng, _bank, kind) |
| img = render(L, rng) |
| name = f"images/{i:06d}.jpg" |
| cv2.imwrite(str(Path(out) / name), img[:, :, ::-1], [cv2.IMWRITE_JPEG_QUALITY, 90]) |
| segs = [{**s, "quad": np.round(s["quad"], 1).tolist()} for s in L.segments()] |
| return {"image": name, "width": L.width, "height": L.height, "source": "synth", "kind": kind, |
| "doc_type": L.doc_type, "segments": segs, "kv_pairs": [list(p) for p in L.kv], "ignore_regions": [], |
| "use": {"detector": True, "recognizer": False, "labels": True, "relations": True, "blocks": L.blocks_known}, |
| "fields": L.fields} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--pages", type=int, default=10_000) |
| ap.add_argument("--out", default="data/synth") |
| ap.add_argument("--fonts", default="fonts") |
| ap.add_argument("--seed", type=int, default=2024) |
| ap.add_argument("--workers", type=int, default=8) |
| args = ap.parse_args() |
| out = Path(args.out) / "train" |
| (out / "images").mkdir(parents=True, exist_ok=True) |
| with Pool(args.workers, initializer=_init, initargs=(args.fonts,)) as pool, open(out / "pages.jsonl", "w") as fh: |
| for k, rec in enumerate(pool.imap(_one, [(i, args.seed, str(out)) for i in range(args.pages)], chunksize=16)): |
| fh.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| if (k + 1) % 1000 == 0: |
| print(f" {k + 1}/{args.pages}", flush=True) |
| print(f"wrote {args.pages} pages to {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
