File size: 2,377 Bytes
4dce182 9d054d4 4dce182 9d054d4 4dce182 9d054d4 4dce182 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """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()
|