| """Training datasets: each mixes real pages (data/real, built by tinydoc.data.convert) with synthetic |
| documents, either pre-rendered (tinydoc.data.prerender) or generated on the fly. Every sample is a pure |
| function of (seed, index). |
| |
| Data locations default to data/real, data/synth and fonts, and can be redirected with the environment |
| variables TINYDOC_REAL_ROOT, TINYDOC_SYNTH_ROOT and TINYDOC_FONT_ROOT (used on Kaggle). |
| """ |
|
|
| import json |
| import os |
| from functools import lru_cache |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from PIL import Image, ImageDraw |
|
|
| from ..charset import encode, normalize |
| from ..features import MAX_SEGMENTS, segment_features |
| from ..geometry import expand_quad, quad_bbox, quad_size, rect_quad, transform_quads |
| from ..models.detector import BLOCK_SHRINK, KERNEL_END_TRIM, KERNEL_HEIGHT_RATIO |
| from ..models.orientation import PATCH |
| from ..models.recognizer import LINE_HEIGHT, WIDTH_STRIDE, X_STRETCH |
| from ..schema import DOC_ID, LABEL_ID, REL_TYPES |
| from ..synth import text as T |
| from ..synth.augment import geometric, photometric |
| from ..synth.documents import KIND_WEIGHTS, build, render |
| from ..synth.fonts import FontBank, coverage |
|
|
| MAX_LINE_WIDTH = 4096 |
| MAX_TRAIN_WIDTH = 1536 |
| MATH_LINE_P = 0.15 |
| WIDTH_BUCKET = 64 |
| SEGMENT_BUCKET = 32 |
|
|
|
|
| def real_root(): |
| return os.environ.get("TINYDOC_REAL_ROOT", "data/real") |
|
|
|
|
| def synth_root(): |
| return os.environ.get("TINYDOC_SYNTH_ROOT", "data/synth") |
|
|
|
|
| def font_root(): |
| return os.environ.get("TINYDOC_FONT_ROOT", "fonts") |
|
|
|
|
| @lru_cache(maxsize=4) |
| def font_bank(root, split): |
| return FontBank(root, "val" if split != "train" else "train") |
|
|
|
|
| @lru_cache(maxsize=32) |
| def load_records(path): |
| p = Path(path) |
| if not p.exists(): |
| return [] |
| return [json.loads(l) for l in open(p)] |
|
|
|
|
| def pick_weighted(rng, weights: dict): |
| names = [k for k, v in weights.items() if v > 0] |
| p = np.array([weights[k] for k in names], np.float64) |
| return names[rng.choice(len(names), p=p / p.sum())] |
|
|
|
|
| |
| |
| AUG = { |
| "synth": ((0.35, 0.40, 0.25), 1.0), |
| "doclaynet": ((0.60, 0.30, 0.10), 0.5), |
| "wildreceipt": ((0.70, 0.30, 0.0), 0.5), |
| } |
| DEFAULT_AUG = ((0.6, 0.3, 0.1), 0.5) |
|
|
|
|
| class PagePool: |
| """Samples un-augmented pages (image, quads, ignore flags, ignore regions) plus how to augment them. |
| |
| "synth" uses pre-rendered pages from TINYDOC_SYNTH_ROOT when they exist for the split, and renders |
| on the fly otherwise. |
| """ |
|
|
| def __init__(self, split, use_key, weights): |
| self.split = split |
| self.real = {} |
| for name in weights: |
| base = Path(synth_root() if name == "synth" else real_root()) / ("" if name == "synth" else name) / split |
| recs = [r for r in load_records(str(base / "pages.jsonl")) if r["use"][use_key]] |
| if recs: |
| self.real[name] = (base, recs) |
| self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real} |
| self._cache = {} |
|
|
| def pick(self, rng): |
| """Choose a page and its augmentation. Returns a dict; the image is loaded separately (load).""" |
| src = pick_weighted(rng, self.weights) |
| geo_p, strength = AUG.get(src, DEFAULT_AUG) |
| geo = str(rng.choice(["none", "rotate", "photo"], p=geo_p)) |
| if src == "synth" and src not in self.real: |
| kind = str(rng.choice(list(KIND_WEIGHTS), p=list(KIND_WEIGHTS.values()))) |
| L = build(rng, font_bank(font_root(), self.split), kind) |
| img = np.ascontiguousarray(render(L, rng)[:, :, ::-1]) |
| segs = L.segments() |
| quads = np.array([s["quad"] for s in segs], np.float32).reshape(-1, 4, 2) |
| info = self._summarise(quads, np.array([detector_ignores(s) for s in segs], bool), [], L.width, L.height, |
| block_rects(segs) if L.blocks_known else None) |
| return {"src": src, "img": img, "info": info, "geo": geo, "strength": strength, |
| "thermal": kind == "receipt" and rng.random() < 0.4} |
| base, recs = self.real[src] |
| i = int(rng.integers(len(recs))) |
| rec = recs[i] |
| return {"src": src, "path": str(base / rec["image"]), "info": self.info(src, i), "geo": geo, "strength": strength, |
| "thermal": rec.get("kind") == "receipt" and rng.random() < 0.4} |
|
|
| @staticmethod |
| def _summarise(quads, ignore, regions, W, H, blocks=None): |
| if len(quads): |
| v = quads[:, 3] - quads[:, 0] |
| heights = np.linalg.norm(v, axis=1) |
| ok = heights[~ignore] if (~ignore).any() else heights |
| med = float(np.median(ok)) |
| bboxes = np.concatenate([quads.min(1), quads.max(1)], 1) |
| else: |
| med, bboxes = 20.0, np.zeros((0, 4), np.float32) |
| |
| return {"quads": quads, "ignore": ignore, "regions": regions, "med_h": med, "bboxes": bboxes, "W": W, "H": H, |
| "blocks": blocks} |
|
|
| def info(self, src, i): |
| """Per-page box arrays and statistics, cached per worker (pages are revisited many times).""" |
| key = (src, i) |
| if key not in self._cache: |
| rec = self.real[src][1][i] |
| quads = np.array([s["quad"] for s in rec["segments"]], np.float32).reshape(-1, 4, 2) |
| ignore = np.array([detector_ignores(s) for s in rec["segments"]], bool) |
| regions = [np.asarray(q, np.float32) for q in rec["ignore_regions"]] |
| blocks = block_rects(rec["segments"]) if rec["use"].get("blocks") else None |
| self._cache[key] = self._summarise(quads, ignore, regions, rec["width"], rec["height"], blocks) |
| return self._cache[key] |
|
|
| @staticmethod |
| def load(pick, reduce=1): |
| """BGR page image, decoded at 1/reduce size (JPEG DCT scaling is much cheaper than resizing).""" |
| if "img" in pick: |
| return pick["img"] |
| flag = {1: cv2.IMREAD_COLOR, 2: cv2.IMREAD_REDUCED_COLOR_2, 4: cv2.IMREAD_REDUCED_COLOR_4}[reduce] |
| return cv2.imread(pick["path"], flag) |
|
|
| def sample(self, rng): |
| """Full page as RGB (orientation data and evaluation).""" |
| p = self.pick(rng) |
| info = p["info"] |
| img = cv2.cvtColor(self.load(p), cv2.COLOR_BGR2RGB) |
| return (img, list(info["quads"]), list(info["ignore"]), list(info["regions"]), p["src"], p["geo"], |
| p["strength"], p["thermal"], info["blocks"]) |
|
|
|
|
| def block_rects(segs): |
| """Layout blocks of a page from its segments' block ids: (rects (B, 4), median line height per block).""" |
| groups = {} |
| for s in segs: |
| if s.get("block") is not None and not s.get("ignore"): |
| groups.setdefault(s["block"], []).append(np.asarray(s["quad"], np.float32)) |
| if not groups: |
| return np.zeros((0, 4), np.float32), np.zeros(0, np.float32) |
| rects, hs = [], [] |
| for qs in groups.values(): |
| q = np.stack(qs) |
| rects.append([q[:, :, 0].min(), q[:, :, 1].min(), q[:, :, 0].max(), q[:, :, 1].max()]) |
| hs.append(float(np.median(q[:, :, 1].max(1) - q[:, :, 1].min(1)))) |
| return np.array(rects, np.float32), np.array(hs, np.float32) |
|
|
|
|
| def block_targets(h, w, block_quads, probe_quads, known, regions): |
| """Block-kernel target and mask. probe_quads: one line-height square per block, transformed with the page, |
| giving the local line height for the shrink.""" |
| kernel = np.zeros((h, w), np.uint8) |
| mask = np.full((h, w), 1 if known else 0, np.uint8) |
| if known: |
| for bq, pq in zip(block_quads, probe_quads): |
| lh = quad_size(pq)[1] |
| d = BLOCK_SHRINK * lh |
| bw, bh = quad_size(bq) |
| if bw > 2 * d + 1 and bh > 2 * d + 1: |
| cv2.fillPoly(kernel, [np.round(expand_quad(bq, -d, -d, -d, -d)).astype(np.int32)], 1) |
| for r in regions: |
| cv2.fillPoly(mask, [np.round(r).astype(np.int32)], 0) |
| return kernel, mask |
|
|
|
|
| def detector_ignores(seg): |
| """Segments the detector neither learns nor is scored on: unreadable text and display maths.""" |
| return bool(seg["ignore"]) or seg.get("label") == "formula" |
|
|
|
|
| def page_region(img, quads, regions, cx, cy, scale, side, fill): |
| """Cut a side x side window centred on (cx, cy) of the page scaled by `scale`; move boxes with it. |
| |
| Cropping before resizing/warping keeps per-sample cost independent of page size. |
| """ |
| H, W = img.shape[:2] |
| src_side = side / scale |
| x0, y0 = cx - src_side / 2, cy - src_side / 2 |
| sx0, sy0 = int(max(0, np.floor(x0))), int(max(0, np.floor(y0))) |
| sx1, sy1 = int(min(W, np.ceil(x0 + src_side))), int(min(H, np.ceil(y0 + src_side))) |
| out = np.empty((side, side, 3), np.uint8) |
| out[:] = fill |
| if sx1 <= sx0 or sy1 <= sy0: |
| return out, [], [] |
| patch = img[sy0:sy1, sx0:sx1] |
| ow, oh = max(1, round((sx1 - sx0) * scale)), max(1, round((sy1 - sy0) * scale)) |
| patch = cv2.resize(patch, (ow, oh), interpolation=cv2.INTER_AREA if scale < 1 else cv2.INTER_LINEAR) |
| ox, oy = int(round((sx0 - x0) * scale)), int(round((sy0 - y0) * scale)) |
| px0, py0 = max(0, ox), max(0, oy) |
| px1, py1 = min(side, ox + ow), min(side, oy + oh) |
| out[py0:py1, px0:px1] = patch[py0 - oy:py1 - oy, px0 - ox:px1 - ox] |
| sx, sy = ow / (sx1 - sx0), oh / (sy1 - sy0) |
| m = np.array([[sx, 0, ox - sx0 * sx], [0, sy, oy - sy0 * sy], [0, 0, 1]]) |
| return out, transform_quads(quads, m), transform_quads(regions, m) |
|
|
|
|
| |
| def kernel_quad(q): |
| w, h = quad_size(q) |
| dy = h * (1 - KERNEL_HEIGHT_RATIO) / 2 |
| dx = min(KERNEL_END_TRIM * h, w * 0.25) |
| return expand_quad(q, -dx, -dy, -dx, -dy) |
|
|
|
|
| def detector_targets(h, w, quads, ignore, regions, min_h=5): |
| kernel = np.zeros((h, w), np.uint8) |
| mask = np.ones((h, w), np.uint8) |
| for q, ign in zip(quads, ignore): |
| _, qh = quad_size(q) |
| if ign or qh < min_h: |
| cv2.fillPoly(mask, [np.round(q).astype(np.int32)], 0) |
| continue |
| cv2.fillPoly(kernel, [np.round(kernel_quad(q)).astype(np.int32)], 1) |
| for r in regions: |
| cv2.fillPoly(mask, [np.round(r).astype(np.int32)], 0) |
| return kernel, mask |
|
|
|
|
| class DetectorDataset(torch.utils.data.Dataset): |
| def __init__(self, split="train", length=100_000, crop=640, seed=0, weights=None): |
| weights = weights or {"synth": 0.45, "doclaynet": 0.3, "wildreceipt": 0.25} |
| self.pool = PagePool(split, "detector", weights) |
| self.length, self.crop, self.seed = length, crop, seed |
|
|
| def __len__(self): |
| return self.length |
|
|
| def __getitem__(self, idx): |
| rng = np.random.default_rng([self.seed, idx]) |
| p = self.pool.pick(rng) |
| info = p["info"] |
| W, H = info["W"], info["H"] |
| scale = np.exp(rng.uniform(np.log(7), np.log(56))) / info["med_h"] |
| scale = float(np.clip(scale, 320 / max(H, W), 3000 / max(H, W))) |
| reduce = 4 if scale <= 0.22 else 2 if scale <= 0.45 else 1 |
| img = self.pool.load(p, reduce) |
| rs = img.shape[1] / W |
| c = self.crop |
| quads, bboxes = info["quads"], info["bboxes"] |
| if len(quads) and rng.random() < 0.75: |
| cx, cy = quads[rng.integers(len(quads))].mean(0) |
| else: |
| cx, cy = rng.uniform(0, W), rng.uniform(0, H) |
| cx += rng.uniform(-c / 4, c / 4) / scale |
| cy += rng.uniform(-c / 4, c / 4) / scale |
| side = int(c * 1.25) |
| half = side / scale |
| near = (bboxes[:, 2] > cx - half) & (bboxes[:, 0] < cx + half) & (bboxes[:, 3] > cy - half) & (bboxes[:, 1] < cy + half) |
| quads_n, ignore = quads[near] * rs, info["ignore"][near] |
| regions = [r * rs for r in info["regions"]] |
| known = info["blocks"] is not None |
| bq = pq = np.zeros((0, 4, 2), np.float32) |
| if known: |
| rects, lhs = info["blocks"] |
| sel = (rects[:, 2] > cx - half) & (rects[:, 0] < cx + half) & (rects[:, 3] > cy - half) & (rects[:, 1] < cy + half) |
| bq = np.array([rect_quad(*r) for r in rects[sel]], np.float32).reshape(-1, 4, 2) * rs |
| pq = np.array([rect_quad(r[0], r[1], r[0] + lh, r[1] + lh) for r, lh in zip(rects[sel], lhs[sel])], |
| np.float32).reshape(-1, 4, 2) * rs |
| nq, nb = len(quads_n), len(bq) |
| allq = np.concatenate([np.asarray(quads_n, np.float32).reshape(-1, 4, 2), bq, pq]) |
| fill = np.median(img[::16, ::16].reshape(-1, 3), 0).astype(np.uint8) |
| region, allq, regions = page_region(img, allq, regions, cx * rs, cy * rs, scale / rs, side, fill) |
| region = np.ascontiguousarray(region[:, :, ::-1]) |
| region, m = geometric(region, rng, p["geo"]) |
| rh, rw = region.shape[:2] |
| x0, y0 = (rw - c) // 2, (rh - c) // 2 |
| m = np.array([[1, 0, -x0], [0, 1, -y0], [0, 0, 1]], np.float64) @ m |
| allq, regions = transform_quads(allq, m), transform_quads(regions, m) |
| quads_n, bq, pq = allq[:nq], allq[nq:nq + nb], allq[nq + nb:] |
| out = region[y0:y0 + c, x0:x0 + c] |
| kernel, mask = detector_targets(c, c, quads_n, ignore, regions) |
| bkernel, bmask = block_targets(c, c, bq, pq, known, regions) |
| out = photometric(np.ascontiguousarray(out), rng, p["strength"], p["thermal"]) |
| return (torch.from_numpy(out).permute(2, 0, 1).float() / 255, |
| torch.from_numpy(np.stack([kernel, bkernel])).float(), torch.from_numpy(np.stack([mask, bmask])).float()) |
|
|
|
|
| |
| def synth_line(rng, bank): |
| """Render one text line with margins. Returns (gray crop, tight box, text).""" |
| for _ in range(10): |
| face, bax = bank.pick(rng, bold=rng.random() < 0.2, italic_ok=True) |
| text = FontBank.supported(face, normalize(T.line_text(rng))).strip() |
| if text: |
| break |
| else: |
| text, (face, bax) = "0", bank.pick(rng) |
| size = int(rng.integers(18, 56)) |
| font = bank.font(face, size, bax) |
| |
| words = text.split(" ") |
| space = font.getlength(" ") * (rng.uniform(0.45, 1.8) if len(words) > 1 and rng.random() < 0.35 else 1.0) |
| placed, x = [], 0.0 |
| for wd in words: |
| if wd: |
| placed.append((x, wd, font, 0.0)) |
| x += font.getlength(wd) + space |
| arr, box = _line_canvas(rng, placed, font) |
| return arr, box, text |
|
|
|
|
| def synth_math_line(rng, bank): |
| """A line with inline maths: sub/superscripts typeset small and raised or lowered (labelled x², H₂O), Greek and |
| maths symbols drawn from a maths font when the text font lacks them. Returns (gray crop, tight box, text).""" |
| for _ in range(10): |
| runs = T.math_runs(T.math_line(rng)) |
| face, bax = bank.pick(rng, bold=rng.random() < 0.1, italic_ok=True) |
| cov = coverage(face.path) |
| if all(c == " " or c in cov or bank.fallback(c) for t, _ in runs for c in t): |
| break |
| else: |
| return synth_line(rng, bank) |
| size = int(rng.integers(20, 56)) |
| font = bank.font(face, size, bax) |
| small_size = max(6, round(size * rng.uniform(0.58, 0.75))) |
| rise, drop = size * rng.uniform(0.28, 0.45), size * rng.uniform(0.10, 0.22) |
| placed, x = [], 0.0 |
| for text, kind in runs: |
| spelled = T.scripted(text, kind) if kind else text |
| if kind and spelled != text and rng.random() < 0.3 and all(c in cov for c in spelled): |
| text, kind = spelled, "" |
| pt, dy = (size, 0.0) if not kind else (small_size, -rise if kind == "^" else drop) |
| for ch in text: |
| if ch == " ": |
| x += font.getlength(" ") |
| continue |
| f = bank.font(face, pt, bax) if ch in cov else bank.font(bank.fallback(ch), pt) |
| placed.append((x, ch, f, dy)) |
| x += f.getlength(ch) |
| if not placed: |
| return synth_line(rng, bank) |
| arr, box = _line_canvas(rng, placed, font) |
| return arr, box, normalize(T.runs_label(runs)) |
|
|
|
|
| def _line_canvas(rng, placed, font): |
| """placed: [(x, text, font, dy)] along one baseline. Draws them on a page-coloured canvas with margins, |
| sometimes with neighbouring lines, rules and a slight tilt. Returns (gray crop, tight box).""" |
| boxes = [(x + bb[0], bb[1] + dy, x + bb[2], bb[3] + dy) |
| for x, s, f, dy in placed for bb in [f.getbbox(s, anchor="ls")]] |
| l, r = min(b[0] for b in boxes), max(b[2] for b in boxes) |
| _, rt, _, rb = font.getbbox("Hg", anchor="ls") |
| top, bot = min(rt, min(b[1] for b in boxes)), max(rb, max(b[3] for b in boxes)) |
| h, w = bot - top, max(r - l, 1) |
| mx, my = 0.6 * h, 0.35 * h |
| W, H = int(w + 2 * mx), int(h + 2 * my) |
| bg = int(rng.integers(170, 256)) |
| ink = int(rng.integers(0, max(bg - 90, 1))) |
| im = Image.new("L", (W, H), bg) |
| d = ImageDraw.Draw(im) |
| for x, s, f, dy in placed: |
| d.text((mx - l + x, my - top + dy), s, font=f, fill=ink, anchor="ls") |
| if rng.random() < 0.3: |
| for sign in (-1, 1): |
| if rng.random() < 0.6: |
| d.text((mx - l + rng.uniform(-h, h), my - top + sign * h * rng.uniform(1.05, 1.5)), |
| T.line_text(rng), font=font, fill=ink, anchor="ls") |
| if rng.random() < 0.06: |
| yy = my + h * rng.uniform(0.85, 1.25) |
| d.line([0, yy, W, yy], fill=ink, width=max(1, int(h / 20))) |
| if rng.random() < 0.06: |
| for xx in (mx * rng.uniform(0, 0.7), W - mx * rng.uniform(0, 0.7)): |
| d.line([xx, 0, xx, H], fill=ink, width=max(1, int(h / 20))) |
| arr = np.array(im) |
| if rng.random() < 0.3: |
| m = cv2.getRotationMatrix2D((W / 2, H / 2), rng.normal(0, 0.8), 1.0) |
| m[0, 1] += rng.uniform(-0.15, 0.15) |
| m[0, 2] -= m[0, 1] * H / 2 |
| arr = cv2.warpAffine(arr, m, (W, H), borderValue=bg) |
| return arr, (mx, my, mx + w, my + h) |
|
|
|
|
| def jitter_crop(gray, box, rng): |
| x0, y0, x1, y1 = box |
| h = y1 - y0 |
| H, W = gray.shape[:2] |
| cx0 = int(np.clip(x0 - rng.uniform(-0.05, 0.5) * h, 0, W - 1)) |
| cx1 = int(np.clip(x1 + rng.uniform(-0.05, 0.5) * h, cx0 + 1, W)) |
| cy0 = int(np.clip(y0 - rng.uniform(-0.08, 0.3) * h, 0, H - 1)) |
| cy1 = int(np.clip(y1 + rng.uniform(-0.08, 0.3) * h, cy0 + 1, H)) |
| return gray[cy0:cy1, cx0:cx1] |
|
|
|
|
| def line_to_uint8(gray) -> np.ndarray: |
| """Grayscale crop -> uint8 (32, W), W a multiple of the CNN stride.""" |
| h, w = gray.shape[:2] |
| nw = int(round(w * LINE_HEIGHT / max(h, 1) * X_STRETCH)) |
| nw = int(np.clip(nw, WIDTH_STRIDE * 2, MAX_LINE_WIDTH)) |
| nw = (nw + WIDTH_STRIDE - 1) // WIDTH_STRIDE * WIDTH_STRIDE |
| return cv2.resize(gray, (nw, LINE_HEIGHT), interpolation=cv2.INTER_AREA if h > LINE_HEIGHT else cv2.INTER_LINEAR) |
|
|
|
|
| def line_to_input(gray) -> np.ndarray: |
| """Grayscale crop -> float32 (32, W) in [0,1].""" |
| return line_to_uint8(gray).astype(np.float32) / 255.0 |
|
|
|
|
| class LineDataset(torch.utils.data.Dataset): |
| def __init__(self, split="train", length=1_000_000, seed=0, weights=None): |
| weights = weights or {"synth": 0.4, "cord-v2": 0.25, "doclaynet": 0.35} |
| self.split, self.length, self.seed = split, length, seed |
| self.real = {} |
| for name in weights: |
| if name != "synth": |
| recs = load_records(str(Path(real_root()) / name / split / "lines.jsonl")) |
| if recs: |
| self.real[name] = (Path(real_root()) / name / split, recs) |
| self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real} |
|
|
| def __len__(self): |
| return self.length |
|
|
| def __getitem__(self, idx): |
| rng = np.random.default_rng([self.seed, idx]) |
| for _ in range(20): |
| out = self._sample(rng) |
| if out[0].shape[-1] <= MAX_TRAIN_WIDTH: |
| return out |
| return out |
|
|
| def _sample(self, rng): |
| src = pick_weighted(rng, self.weights) |
| tag = src |
| if src == "synth": |
| bank = font_bank(font_root(), self.split) |
| maths = rng.random() < MATH_LINE_P |
| gray, box, text = (synth_math_line if maths else synth_line)(rng, bank) |
| strength, tag = 1.0, "synth-math" if maths else "synth" |
| else: |
| base, recs = self.real[src] |
| r = recs[rng.integers(len(recs))] |
| gray = cv2.imread(str(base / r["path"]), cv2.IMREAD_GRAYSCALE) |
| box, text, strength = r["box"], r["text"], 0.5 |
| crop = jitter_crop(gray, box, rng) |
| if rng.random() < 0.2: |
| |
| th = max(box[3] - box[1], 1) |
| f = min(1.0, rng.uniform(7, 13) / th) |
| if f < 0.95: |
| h0, w0 = crop.shape[:2] |
| small = cv2.resize(crop, (max(1, round(w0 * f)), max(1, round(h0 * f))), interpolation=cv2.INTER_AREA) |
| crop = cv2.resize(small, (w0, h0), interpolation=cv2.INTER_LINEAR if rng.random() < 0.5 else cv2.INTER_CUBIC) |
| crop = photometric(crop, rng, strength, thermal=src == "synth" and rng.random() < 0.1) |
| if rng.random() < 0.08: |
| crop = 255 - crop |
| x = line_to_uint8(crop) |
| return torch.from_numpy(x)[None], torch.tensor(encode(text), dtype=torch.long), tag |
|
|
|
|
| def collate_lines(batch): |
| imgs, targets, srcs = zip(*batch) |
| wmax = max(i.shape[-1] for i in imgs) |
| wmax = (wmax + WIDTH_BUCKET - 1) // WIDTH_BUCKET * WIDTH_BUCKET |
| B = len(imgs) |
| x = torch.ones(B, 1, LINE_HEIGHT, wmax) |
| steps = torch.tensor([i.shape[-1] // WIDTH_STRIDE for i in imgs]) |
| for k, im in enumerate(imgs): |
| x[k, :, :, : im.shape[-1]] = im.float() / 255.0 |
| step_mask = torch.arange(wmax // WIDTH_STRIDE)[None, :] >= steps[:, None] |
| return x, step_mask, steps, torch.cat(targets), torch.tensor([len(t) for t in targets]), list(srcs), targets |
|
|
|
|
| |
| def text_patches(gray, n, rng=None): |
| """Pick up to n PATCH-sized windows dense in ink-like edges. Returns list of (x, y).""" |
| H, W = gray.shape |
| if H < PATCH or W < PATCH: |
| return [(0, 0)] |
| ink = cv2.adaptiveThreshold(gray, 1, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 15).astype(np.float32) |
| s = 8 |
| small = cv2.resize(ink, (W // s, H // s), interpolation=cv2.INTER_AREA) |
| k = PATCH // s |
| dens = cv2.boxFilter(small, -1, (k, k), normalize=True, borderType=cv2.BORDER_CONSTANT) |
| dens[: k // 2, :] = dens[-(k // 2):, :] = 0 |
| dens[:, : k // 2] = dens[:, -(k // 2):] = 0 |
| |
| score = dens * (dens < 0.35) |
| out = [] |
| for _ in range(n): |
| if rng is not None: |
| flat = score.ravel() ** 2 |
| if flat.sum() <= 0: |
| break |
| i = rng.choice(len(flat), p=flat / flat.sum()) |
| else: |
| i = int(score.argmax()) |
| if score.ravel()[i] <= 0.01: |
| break |
| cy, cx = divmod(i, score.shape[1]) |
| out.append((int(np.clip(cx * s - PATCH // 2, 0, W - PATCH)), int(np.clip(cy * s - PATCH // 2, 0, H - PATCH)))) |
| score[max(0, cy - k // 2): cy + k // 2, max(0, cx - k // 2): cx + k // 2] = 0 |
| return out or [((W - PATCH) // 2, (H - PATCH) // 2)] |
|
|
|
|
| ORIENT_LONG_SIDE = 1280 |
|
|
|
|
| def orientation_page(img_rgb, long_side=ORIENT_LONG_SIDE): |
| gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY) if img_rgb.ndim == 3 else img_rgb |
| s = long_side / max(gray.shape) |
| return cv2.resize(gray, (max(1, round(gray.shape[1] * s)), max(1, round(gray.shape[0] * s))), interpolation=cv2.INTER_AREA) |
|
|
|
|
| class OrientationDataset(torch.utils.data.Dataset): |
| """Each item is `per_page` patches from one page, each independently rotated by k*90 degrees.""" |
|
|
| def __init__(self, split="train", length=200_000, seed=0, weights=None, per_page=8): |
| weights = weights or {"synth": 0.5, "doclaynet": 0.25, "wildreceipt": 0.25} |
| self.pool = PagePool(split, "detector", weights) |
| self.length, self.seed, self.per_page = length, seed, per_page |
|
|
| def __len__(self): |
| return self.length |
|
|
| def __getitem__(self, idx): |
| rng = np.random.default_rng([self.seed, idx]) |
| img, quads, ignore, _, _, _, strength, thermal, _ = self.pool.sample(rng) |
| long_side = int(rng.integers(1000, 1600)) |
| gray = orientation_page(img, long_side) |
| scale = long_side / max(img.shape[:2]) |
| centers = [q.mean(0) * scale for q, ig in zip(quads, ignore) if not ig] |
| pos = text_patches(gray, self.per_page, rng) |
| big = int(PATCH * 1.3) |
| padded = cv2.copyMakeBorder(gray, big, big, big, big, cv2.BORDER_REPLICATE) |
| patches, labels = [], [] |
| for i in range(self.per_page): |
| |
| if centers and rng.random() < 0.75: |
| cx, cy = centers[rng.integers(len(centers))] + rng.uniform(-PATCH / 3, PATCH / 3, 2) |
| else: |
| x, y = pos[i % len(pos)] |
| cx, cy = x + PATCH / 2, y + PATCH / 2 |
| cx, cy = int(np.clip(cx, 0, gray.shape[1])) + big, int(np.clip(cy, 0, gray.shape[0])) + big |
| win = padded[cy - big // 2: cy - big // 2 + big, cx - big // 2: cx - big // 2 + big] |
| if rng.random() < 0.6: |
| m = cv2.getRotationMatrix2D((big / 2, big / 2), rng.normal(0, 3), 1.0) |
| win = cv2.warpAffine(win, m, (big, big), borderMode=cv2.BORDER_REPLICATE) |
| o = (big - PATCH) // 2 |
| patch = win[o:o + PATCH, o:o + PATCH] |
| k = int(rng.integers(4)) |
| patch = np.ascontiguousarray(np.rot90(patch, -k)) |
| patches.append(photometric(patch, rng, strength, thermal)) |
| labels.append(k) |
| x = torch.from_numpy(np.stack(patches).astype(np.float32) / 255)[:, None] |
| return x, torch.tensor(labels) |
|
|
|
|
| |
| CONFUSIONS = {"0": "O", "O": "0", "o": "0", "1": "l", "l": "1", "I": "1", "5": "S", "S": "5", "8": "B", "B": "8", |
| ",": ".", ".": ",", "2": "Z", "6": "G", "g": "9", "9": "g", "e": "c", "c": "e", "u": "v", ":": ";"} |
|
|
|
|
| def ocr_noise(text, rng, rate): |
| out = [] |
| for ch in text: |
| r = rng.random() |
| if r < rate * 0.6: |
| out.append(CONFUSIONS.get(ch, ch if rng.random() < 0.5 else T.pick(rng, "abcdefghijklmnopqrstuvwxyz0123456789"))) |
| elif r < rate * 0.8: |
| continue |
| elif r < rate: |
| out.extend([ch, T.pick(rng, ".,'-il")]) |
| else: |
| out.append(ch) |
| return "".join(out) or text |
|
|
|
|
| EXTRACTOR_KINDS = {"receipt": 0.3, "invoice": 0.35, "general": 0.1, "newspaper": 0.07, "magazine": 0.06, "academic": 0.12} |
|
|
|
|
| class LayoutDataset(torch.utils.data.Dataset): |
| def __init__(self, split="train", length=200_000, seed=0, weights=None): |
| weights = weights or {"synth": 0.45, "cord-v2": 0.17, "wildreceipt": 0.13, "doclaynet": 0.25} |
| self.split, self.length, self.seed = split, length, seed |
| self.real = {} |
| for name in weights: |
| if name != "synth": |
| recs = [r for r in load_records(str(Path(real_root()) / name / split / "pages.jsonl")) if r["use"]["labels"]] |
| if recs: |
| self.real[name] = recs |
| self.prerendered = load_records(str(Path(synth_root()) / split / "pages.jsonl")) |
| self.weights = {k: v for k, v in weights.items() if k == "synth" or k in self.real} |
|
|
| def __len__(self): |
| return self.length |
|
|
| def sample_doc(self, rng): |
| src = pick_weighted(rng, self.weights) |
| |
| |
| if src == "synth" and (not self.prerendered or rng.random() < 0.5): |
| kind = str(rng.choice(list(EXTRACTOR_KINDS), p=list(EXTRACTOR_KINDS.values()))) |
| L = build(rng, font_bank(font_root(), self.split), kind) |
| segs = [{"bbox": list(s.box), "text": s.text, "label": s.label, "row": s.row, "block": s.block} for s in L.segs] |
| return segs, L.kv, L.doc_type, L.width, L.height, (True, True, L.blocks_known), f"synth-{L.doc_type}" |
| recs = self.prerendered if src == "synth" else self.real[src] |
| rec = recs[rng.integers(len(recs))] |
| segs = [{"bbox": quad_bbox(s["quad"]), "text": s["text"], "label": s["label"], "row": s["row"], "block": s.get("block")} |
| for s in rec["segments"] if not s["ignore"]] |
| keep = [i for i, s in enumerate(rec["segments"]) if not s["ignore"]] |
| remap = {old: new for new, old in enumerate(keep)} |
| kv = [(remap[k], remap[v]) for k, v in rec["kv_pairs"] if k in remap and v in remap] |
| use = rec["use"] |
| label = f"synth-{rec['doc_type']}" if src == "synth" else src |
| return segs, kv, rec["doc_type"], rec["width"], rec["height"], (use["relations"], use["relations"], use.get("blocks", False)), label |
|
|
| def __getitem__(self, idx): |
| rng = np.random.default_rng([self.seed, idx]) |
| segs, kv, doc_type, W, H, rel_known, src = self.sample_doc(rng) |
| if self.split == "train": |
| rate = rng.uniform(0, 0.08) |
| drop_spaces = rng.random() < 0.08 |
| for s in segs: |
| if rng.random() < 0.5: |
| s["text"] = ocr_noise(s["text"], rng, rate) |
| if drop_spaces: |
| s["text"] = s["text"].replace(" ", "") |
| h = s["bbox"][3] - s["bbox"][1] |
| s["bbox"] = [v + rng.normal(0, 0.06 * h) for v in s["bbox"]] |
| s["conf"] = float(np.clip(1 - rate * rng.uniform(2, 8), 0.3, 1)) |
| if rng.random() < 0.3: |
| f = rng.uniform(1.0, 1.8) |
| ox, oy = rng.uniform(0, W * (f - 1)), rng.uniform(0, H * (f - 1)) |
| for s in segs: |
| s["bbox"] = [s["bbox"][0] + ox, s["bbox"][1] + oy, s["bbox"][2] + ox, s["bbox"][3] + oy] |
| W, H = W * f, H * f |
| order = sorted(range(len(segs)), key=lambda i: (round(segs[i]["bbox"][1] / 8), segs[i]["bbox"][0])) |
| if len(order) > MAX_SEGMENTS: |
| start = int(rng.integers(0, len(order) - MAX_SEGMENTS + 1)) |
| order = order[start:start + MAX_SEGMENTS] |
| pos = {old: new for new, old in enumerate(order)} |
| segs = [segs[i] for i in order] |
| n = len(segs) |
| f = segment_features(segs, W, H) |
| labels = np.array([LABEL_ID[s["label"]] if s["label"] is not None else -100 for s in segs], np.int64) |
| rel = np.zeros((len(REL_TYPES), n, n), np.float32) |
| rows = np.array([s["row"] if s["row"] is not None else -1 for s in segs]) |
| same = (rows[:, None] == rows[None, :]) & (rows[:, None] >= 0) |
| np.fill_diagonal(same, False) |
| rel[0] = same |
| for k, v in kv: |
| if k in pos and v in pos: |
| rel[1, pos[k], pos[v]] = 1 |
| blocks = np.array([s["block"] if s.get("block") is not None else -1 for s in segs]) |
| same = (blocks[:, None] == blocks[None, :]) & (blocks[:, None] >= 0) |
| np.fill_diagonal(same, False) |
| rel[2] = same |
| rel_mask = np.array(rel_known, np.float32) |
| return {**{k: torch.from_numpy(v) for k, v in f.items()}, "labels": torch.from_numpy(labels), |
| "rel": torch.from_numpy(rel), "rel_known": torch.from_numpy(rel_mask), |
| "doc": DOC_ID[doc_type], "src": src} |
|
|
|
|
| def collate_layouts(batch): |
| B = len(batch) |
| n = max(max(b["feats"].shape[0] for b in batch), 1) |
| n = (n + SEGMENT_BUCKET - 1) // SEGMENT_BUCKET * SEGMENT_BUCKET |
| L = batch[0]["char_ids"].shape[1] |
| F_ = batch[0]["feats"].shape[1] |
| R = batch[0]["rel"].shape[0] |
| out = {"char_ids": torch.zeros(B, n, L, dtype=torch.long), "feats": torch.zeros(B, n, F_), |
| "rel_boxes": torch.zeros(B, n, 4), "valid": torch.zeros(B, n), "labels": torch.full((B, n), -100, dtype=torch.long), |
| "rel": torch.zeros(B, R, n, n), "rel_known": torch.stack([b["rel_known"] for b in batch]), |
| "doc": torch.tensor([b["doc"] for b in batch]), "src": [b["src"] for b in batch]} |
| for i, b in enumerate(batch): |
| k = b["feats"].shape[0] |
| out["char_ids"][i, :k] = b["char_ids"] |
| out["feats"][i, :k] = b["feats"] |
| out["rel_boxes"][i, :k] = b["rel_boxes"] |
| out["valid"][i, :k] = 1 |
| out["labels"][i, :k] = b["labels"] |
| out["rel"][i, :, :k, :k] = b["rel"] |
| return out |
|
|
