| """Torch datasets and CenterNet-style target encoding.""" |
| from __future__ import annotations |
|
|
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from .config import OUT_STRIDE, PREP, SYNTH, TILE, TRAIN_CSV |
| from .preprocess import load_prepared, stream_tensors |
| from .tiles import boxes_in_tile, pad_to |
|
|
|
|
| def load_gt(csv_path: Path | None = None) -> dict[int, np.ndarray]: |
| """Ground-truth boxes keyed by train pair index.""" |
| df = pd.read_csv(csv_path or TRAIN_CSV) |
| out: dict[int, list] = defaultdict(list) |
| for r in df.itertuples(index=False): |
| idx = int(str(r.template_image).split("_")[-1].split(".")[0]) |
| out[idx].append([r.left_x, r.top_y, r.right_x, r.bottom_y]) |
| return {k: np.asarray(v, np.float32) for k, v in out.items()} |
|
|
|
|
| def gaussian_radius(h: float, w: float, min_overlap: float = 0.5) -> float: |
| """CenterNet radius: keeps a shifted box above min_overlap IoU.""" |
| a1, b1, c1 = 1, (h + w), w * h * (1 - min_overlap) / (1 + min_overlap) |
| r1 = (b1 - np.sqrt(max(b1 ** 2 - 4 * a1 * c1, 0))) / (2 * a1) |
| a2, b2, c2 = 4, 2 * (h + w), (1 - min_overlap) * w * h |
| r2 = (b2 - np.sqrt(max(b2 ** 2 - 4 * a2 * c2, 0))) / (2 * a2) |
| a3 = 4 * min_overlap |
| b3 = -2 * min_overlap * (h + w) |
| c3 = (min_overlap - 1) * w * h |
| r3 = (-b3 + np.sqrt(max(b3 ** 2 - 4 * a3 * c3, 0))) / (2 * a3) |
| return max(1.0, float(min(r1, r2, r3))) |
|
|
|
|
| def draw_gaussian(heatmap: np.ndarray, cx: int, cy: int, radius: int) -> None: |
| diameter = 2 * radius + 1 |
| sigma = diameter / 6.0 |
| ax = np.arange(-radius, radius + 1, dtype=np.float32) |
| g = np.exp(-(ax[None, :] ** 2 + ax[:, None] ** 2) / (2 * sigma ** 2)) |
| h, w = heatmap.shape |
| left, right = min(cx, radius), min(w - cx, radius + 1) |
| top, bottom = min(cy, radius), min(h - cy, radius + 1) |
| if right <= -left or bottom <= -top: |
| return |
| region = heatmap[cy - top:cy + bottom, cx - left:cx + right] |
| masked = g[radius - top:radius + bottom, radius - left:radius + right] |
| np.maximum(region, masked, out=region) |
|
|
|
|
| def encode_targets(boxes: np.ndarray, tile: int = TILE, stride: int = OUT_STRIDE): |
| """Boxes in tile coordinates -> heatmap, wh, offset, regression mask, seg mask.""" |
| size = tile // stride |
| hm = np.zeros((1, size, size), np.float32) |
| wh = np.zeros((2, size, size), np.float32) |
| off = np.zeros((2, size, size), np.float32) |
| reg = np.zeros((1, size, size), np.float32) |
| seg = np.zeros((1, size, size), np.float32) |
|
|
| for x1, y1, x2, y2 in boxes: |
| bw, bh = x2 - x1, y2 - y1 |
| if bw <= 1 or bh <= 1: |
| continue |
| cx, cy = (x1 + x2) / 2 / stride, (y1 + y2) / 2 / stride |
| ix, iy = int(cx), int(cy) |
| if not (0 <= ix < size and 0 <= iy < size): |
| continue |
| draw_gaussian(hm[0], ix, iy, int(gaussian_radius(bh / stride, bw / stride))) |
| wh[:, iy, ix] = [bw, bh] |
| off[:, iy, ix] = [cx - ix, cy - iy] |
| reg[0, iy, ix] = 1.0 |
| sx1, sy1 = int(np.floor(x1 / stride)), int(np.floor(y1 / stride)) |
| sx2, sy2 = int(np.ceil(x2 / stride)), int(np.ceil(y2 / stride)) |
| seg[0, max(0, sy1):min(size, sy2), max(0, sx1):min(size, sx2)] = 1.0 |
| return hm, wh, off, reg, seg |
|
|
|
|
| def photometric_jitter(photo: np.ndarray, rng: np.random.RandomState) -> np.ndarray: |
| """Extra print/scan-like noise on the photo stream only.""" |
| out = photo.astype(np.float32) |
| out = out * rng.uniform(0.92, 1.08) + rng.uniform(-12, 12) |
| if rng.rand() < 0.5: |
| out = cv2.GaussianBlur(out, (0, 0), rng.uniform(0.3, 1.0)) |
| if rng.rand() < 0.7: |
| out += rng.normal(0, rng.uniform(1.0, 5.0), out.shape).astype(np.float32) |
| return np.clip(out, 0, 255).astype(np.uint8) |
|
|
|
|
| class GroupedSampler(torch.utils.data.Sampler): |
| """Shuffle pairs, keep a pair's samples adjacent. |
| |
| Every sample decodes four full-page PNGs, so random access would decode the |
| same pair `samples_per_pair` times per epoch. Grouping lets the one-entry |
| cache in the dataset serve the rest of the group. |
| """ |
|
|
| def __init__(self, n_pairs: int, samples_per_pair: int, seed: int = 0): |
| self.n_pairs = n_pairs |
| self.samples_per_pair = samples_per_pair |
| self.seed = seed |
| self.epoch = 0 |
|
|
| def set_epoch(self, epoch: int) -> None: |
| self.epoch = epoch |
|
|
| def __len__(self) -> int: |
| return self.n_pairs * self.samples_per_pair |
|
|
| def __iter__(self): |
| rng = np.random.RandomState(self.seed + self.epoch) |
| for pair in rng.permutation(self.n_pairs): |
| base = int(pair) * self.samples_per_pair |
| for k in rng.permutation(self.samples_per_pair): |
| yield base + int(k) |
|
|
|
|
| class PairTileDataset(Dataset): |
| """Random tiles from prepared pairs, biased towards annotated differences.""" |
|
|
| def __init__(self, items: list[tuple[str, int]], gt: dict, samples_per_pair: int = 8, |
| tile: int = TILE, positive_ratio: float = 0.6, augment: bool = True, |
| prep_root: Path | None = None): |
| self.items = items |
| self.gt = gt |
| self.samples_per_pair = samples_per_pair |
| self.tile = tile |
| self.positive_ratio = positive_ratio |
| self.augment = augment |
| self.prep_root = prep_root or PREP |
| self._cache_key = None |
| self._cache_val = None |
|
|
| def __len__(self) -> int: |
| return len(self.items) * self.samples_per_pair |
|
|
| def _key(self, split: str, idx: int): |
| return (split, idx) |
|
|
| def _load(self, split: str, idx: int): |
| key = (split, idx) |
| if key != self._cache_key: |
| root = self.prep_root if split != "synth" else SYNTH |
| self._cache_val = load_prepared(split, idx, root) |
| self._cache_key = key |
| return self._cache_val |
|
|
| def __getitem__(self, i: int): |
| split, idx = self.items[i // self.samples_per_pair] |
| rng = np.random.RandomState((i * 9973 + idx * 7919) % (2 ** 31)) |
| t, p, tn, pn = self._load(split, idx) |
| boxes = self.gt.get(self._key(split, idx), np.zeros((0, 4), np.float32)) |
|
|
| t, p = pad_to(t, self.tile), pad_to(p, self.tile) |
| tn, pn = pad_to(tn, self.tile), pad_to(pn, self.tile) |
| h, w = t.shape[:2] |
|
|
| if len(boxes) and rng.rand() < self.positive_ratio: |
| bx = boxes[rng.randint(len(boxes))] |
| cx, cy = (bx[0] + bx[2]) / 2, (bx[1] + bx[3]) / 2 |
| x0 = int(np.clip(cx - self.tile / 2 + rng.randint(-200, 201), 0, max(0, w - self.tile))) |
| y0 = int(np.clip(cy - self.tile / 2 + rng.randint(-200, 201), 0, max(0, h - self.tile))) |
| else: |
| x0 = rng.randint(0, max(1, w - self.tile + 1)) |
| y0 = rng.randint(0, max(1, h - self.tile + 1)) |
|
|
| sl = (slice(y0, y0 + self.tile), slice(x0, x0 + self.tile)) |
| tc, pc, tnc, pnc = t[sl], p[sl], tn[sl], pn[sl] |
| if self.augment: |
| pc = photometric_jitter(pc, rng) |
|
|
| a, b = stream_tensors(tc, pc, tnc, pnc) |
| local = boxes_in_tile(boxes, x0, y0, self.tile) |
| hm, wh, off, reg, seg = encode_targets(local, self.tile) |
| return { |
| "a": torch.from_numpy(a), |
| "b": torch.from_numpy(b), |
| "hm": torch.from_numpy(hm), |
| "wh": torch.from_numpy(wh), |
| "off": torch.from_numpy(off), |
| "reg": torch.from_numpy(reg), |
| "seg": torch.from_numpy(seg), |
| } |
|
|
