| """Data for the bold-line intersection detector (CenterNet style) on v5 canvases. |
| |
| Ground truth: bold 5 mm lines lie at U = ou + 5i, V = ov + 5j; with a = (U - ou) / 5, b = (V - ov) / 5 |
| from the dense uv map, an intersection is found in every 2 x 2 pixel cell by one Newton step. Nodes |
| inside the grid area are kept; a border lying on a bold line keeps its nodes, at other borders nodes |
| closer than MARGIN_TOL_MM are dropped. |
| Windows: mostly >= 44 of 49 patches on the sheet; EDGE_WINDOW_PROB of them show the sheet border. |
| Occlusions: 0-2 blobs (blot, blur, noise, glare); the targets stay unchanged. |
| Targets: 28 x 28 cells of 8 px: heat (Gaussian around each node), offset inside the cell, node flag. |
| """ |
|
|
| import math |
|
|
| import cv2 |
| import numpy as np |
|
|
| import canvas |
| import crops |
|
|
| CROP_SIZE = crops.CROP_SIZE |
| CELL = 8 |
| CELLS = CROP_SIZE // CELL |
| HEAT_SIGMA = 0.7 |
| EDGE_WINDOW_PROB = 0.3 |
| EDGE_MIN_VALID = 20 |
| MARGIN_TOL_MM = 0.3 |
| OCCLUSION_COUNT_P = (0.2, 0.5, 0.3) |
| OCCLUSION_AREA_PATCHES = (0.6, 2.2) |
| OCCLUSION_KINDS = ("blot", "blur", "noise", "glare") |
| PLAIN_KEEP_EVERY = 3 |
| KIND_CODES = {"red": 0, "dots": 1, "plain_1mm": 2} |
| CH_HEAT, CH_OFF_X, CH_OFF_Y, CH_POS = range(4) |
|
|
|
|
| def canvas_points(uv_mm, valid, texture_info): |
| """Bold-line intersections in a canvas. -> [K, 2] x, y (canvas px).""" |
| grid = texture_info["grid"] |
| if grid["kind"] == "plain_1mm": |
| return np.zeros((0, 2)) |
| ou, ov = grid["bold_line_mm"] |
| a = (uv_mm[..., 0].astype(np.float64) - ou) / 5.0 |
| b = (uv_mm[..., 1].astype(np.float64) - ov) / 5.0 |
| ok = valid[:-1, :-1] & valid[:-1, 1:] & valid[1:, :-1] & valid[1:, 1:] |
| ac = 0.25 * (a[:-1, :-1] + a[:-1, 1:] + a[1:, :-1] + a[1:, 1:]) |
| bc = 0.25 * (b[:-1, :-1] + b[:-1, 1:] + b[1:, :-1] + b[1:, 1:]) |
| ax = 0.5 * (a[:-1, 1:] - a[:-1, :-1] + a[1:, 1:] - a[1:, :-1]) |
| ay = 0.5 * (a[1:, :-1] - a[:-1, :-1] + a[1:, 1:] - a[:-1, 1:]) |
| bx = 0.5 * (b[:-1, 1:] - b[:-1, :-1] + b[1:, 1:] - b[1:, :-1]) |
| by = 0.5 * (b[1:, :-1] - b[:-1, :-1] + b[1:, 1:] - b[:-1, 1:]) |
| det = ax * by - ay * bx |
| i, j = np.round(ac), np.round(bc) |
| da, db = i - ac, j - bc |
| with np.errstate(divide="ignore", invalid="ignore"): |
| dx = (by * da - ay * db) / det |
| dy = (-bx * da + ax * db) / det |
| hit = ok & (np.abs(det) > 1e-9) & (np.abs(dx) <= 0.5) & (np.abs(dy) <= 0.5) |
| ys, xs = np.nonzero(hit) |
| if len(ys) == 0: |
| return np.zeros((0, 2)) |
| px = xs + 0.5 + dx[ys, xs] |
| py = ys + 0.5 + dy[ys, xs] |
| ii, jj = i[ys, xs], j[ys, xs] |
| u, v = ou + 5.0 * ii, ov + 5.0 * jj |
| x0, y0, x1, y1 = grid["area_mm"] |
| bold = grid["edges_bold"] |
| width_mm, height_mm = texture_info["size_mm"] |
| t = MARGIN_TOL_MM |
| lo_u = max(x0 - t if bold[0] else x0 + t, t) |
| lo_v = max(y0 - t if bold[1] else y0 + t, t) |
| hi_u = min(x1 + t if bold[2] else x1 - t, width_mm - t) |
| hi_v = min(y1 + t if bold[3] else y1 - t, height_mm - t) |
| inside = (u > lo_u) & (u < hi_u) & (v > lo_v) & (v < hi_v) |
| keys = np.stack([ii, jj], axis=1)[inside] |
| if len(keys) == 0: |
| return np.zeros((0, 2)) |
| |
| _, group = np.unique(keys, axis=0, return_inverse=True) |
| group = group.reshape(-1) |
| sums = np.zeros((group.max() + 1, 2)) |
| np.add.at(sums, group, np.stack([px, py], axis=1)[inside]) |
| return sums / np.bincount(group)[:, None] |
|
|
|
|
| def points_to_crop(points, window): |
| x0, y0, s = window |
| x = (points[:, 0] + 0.5 - x0) / s - 0.5 |
| y = (points[:, 1] + 0.5 - y0) / s - 0.5 |
| keep = (x >= -0.5) & (x < CROP_SIZE - 0.5) & (y >= -0.5) & (y < CROP_SIZE - 0.5) |
| return np.stack([x, y], axis=1)[keep] |
|
|
|
|
| def _blob_mask(rng, size): |
| """Soft rotated ellipse with a wobbly border, 0..1.""" |
| area = rng.uniform(*OCCLUSION_AREA_PATCHES) * crops.PATCH_SIZE ** 2 |
| aspect = rng.uniform(1.0, 2.5) |
| ra = math.sqrt(area * aspect / math.pi) |
| rb = ra / aspect |
| cx, cy = rng.uniform(0.1, 0.9) * size, rng.uniform(0.1, 0.9) * size |
| angle = rng.uniform(0, math.pi) |
| soft = rng.uniform(0.8, 3.0) |
| reach = int(math.ceil(1.4 * ra + 3 * soft + 2)) |
| xa, xb = max(0, int(cx) - reach), min(size, int(cx) + reach + 1) |
| ya, yb = max(0, int(cy) - reach), min(size, int(cy) + reach + 1) |
| yy, xx = np.mgrid[ya:yb, xa:xb].astype(np.float32) |
| xr = (xx - cx) * math.cos(angle) + (yy - cy) * math.sin(angle) |
| yr = -(xx - cx) * math.sin(angle) + (yy - cy) * math.cos(angle) |
| radius = np.sqrt((xr / ra) ** 2 + (yr / rb) ** 2) |
| coarse = rng.normal(0, 1, (6, 6)).astype(np.float32) |
| wobble = cv2.resize(coarse, (xb - xa, yb - ya), interpolation=cv2.INTER_CUBIC) * rng.uniform(0.1, 0.3) |
| mask = np.zeros((size, size), np.float32) |
| mask[ya:yb, xa:xb] = (radius + wobble < 1.0).astype(np.float32) |
| return cv2.GaussianBlur(mask, (0, 0), soft) |
|
|
|
|
| def _smooth_field(rng, size, cells, amplitude): |
| coarse = rng.normal(0, 1, (cells, cells)).astype(np.float32) |
| return cv2.resize(coarse, (size, size), interpolation=cv2.INTER_CUBIC) * amplitude |
|
|
|
|
| def occlude(rng, crop_bgr): |
| """0-2 occlusion blobs.""" |
| size = crop_bgr.shape[0] |
| count = int(rng.choice(len(OCCLUSION_COUNT_P), p=OCCLUSION_COUNT_P)) |
| image = crop_bgr.astype(np.float32) |
| for _ in range(count): |
| alpha = _blob_mask(rng, size) |
| kind = OCCLUSION_KINDS[int(rng.integers(len(OCCLUSION_KINDS)))] |
| if kind == "blot": |
| color = np.array([[150, 60, 25], [30, 30, 35], [45, 85, 140], [60, 60, 170], |
| [20, 50, 90]][int(rng.integers(5))], np.float32) |
| color = color * rng.uniform(0.7, 1.3) |
| layer = color[None, None, :] * (1.0 + _smooth_field(rng, size, 12, 0.15))[..., None] |
| layer += rng.normal(0, 6, (size, size, 1)).astype(np.float32) |
| strength = rng.uniform(0.8, 1.0) |
| elif kind == "blur": |
| layer = cv2.GaussianBlur(image, (0, 0), rng.uniform(5.0, 12.0)) |
| strength = 1.0 |
| elif kind == "noise": |
| sigma = rng.uniform(45.0, 100.0) |
| grain = rng.normal(0, sigma, (size, size, 3 if rng.random() < 0.5 else 1)).astype(np.float32) |
| layer = cv2.GaussianBlur(image, (0, 0), 1.5) + grain |
| strength = 1.0 |
| else: |
| level = rng.uniform(235, 255) |
| tint = np.array([1.0, rng.uniform(0.97, 1.0), rng.uniform(0.95, 1.0)], np.float32) |
| layer = np.full_like(image, level) * tint |
| alpha = np.clip(cv2.GaussianBlur(alpha, (0, 0), rng.uniform(3.0, 8.0)) * 1.4, 0, 1) |
| strength = rng.uniform(0.9, 1.0) |
| a = (alpha * strength)[..., None] |
| image = image * (1 - a) + layer * a |
| return np.clip(image, 0, 255).astype(np.uint8) |
|
|
|
|
| def make_targets(points_crop): |
| """-> [28, 28, 4] float32: heat, offset x, offset y, node flag.""" |
| target = np.zeros((CELLS, CELLS, 4), np.float32) |
| if len(points_crop) == 0: |
| return target |
| pc = (points_crop + 0.5) / CELL |
| cell = np.clip(np.floor(pc).astype(int), 0, CELLS - 1) |
| offset = np.clip(pc - cell, 0.0, 1.0 - 1e-6) |
| grid = np.arange(CELLS) |
| heat = np.zeros((CELLS, CELLS), np.float32) |
| for (cx, cy) in cell: |
| heat = np.maximum(heat, np.exp(-((grid[None, :] - cx) ** 2 + (grid[:, None] - cy) ** 2) / (2 * HEAT_SIGMA ** 2))) |
| for (cx, cy), (ox, oy) in zip(cell, offset): |
| heat[cy, cx] = 1.0 |
| target[cy, cx, CH_OFF_X] = ox |
| target[cy, cx, CH_OFF_Y] = oy |
| target[cy, cx, CH_POS] = 1.0 |
| target[..., CH_HEAT] = heat |
| return target |
|
|
|
|
| def train_selection(directory, indices): |
| """No "dots" sheets (their dots are invisible at 3.5-7.5 px/mm), every PLAIN_KEEP_EVERY-th "plain_1mm" one.""" |
| kind = {i: canvas.read_texture_info(directory, i)["grid"]["kind"] for i in indices} |
| keep_plain = set([i for i in indices if kind[i] == "plain_1mm"][::PLAIN_KEEP_EVERY]) |
| return [i for i in indices if kind[i] == "red" or i in keep_plain] |
|
|
|
|
| def make_sample(rng, directory, index, cache): |
| """-> (crop BGR uint8, target [28, 28, 4], grid kind code).""" |
| image, uv_mm, valid = canvas.read_sample(directory, index) |
| texture = canvas.read_texture_info(directory, index) |
| if rng.random() < EDGE_WINDOW_PROB: |
| window = crops.choose_window(rng, uv_mm, valid, image.shape[0], EDGE_MIN_VALID, crops.MIN_VALID_PATCHES) |
| else: |
| window = crops.choose_window(rng, uv_mm, valid, image.shape[0]) |
| crop = occlude(rng, crops.crop_image(image, window)) |
| if index not in cache: |
| cache[index] = canvas_points(uv_mm, valid, texture) |
| return crop, make_targets(points_to_crop(cache[index], window)), KIND_CODES[texture["grid"]["kind"]] |
|
|
|
|
| class PointDataset: |
| """train: new random window and occlusions every call; otherwise one seeded window per index.""" |
|
|
| def __init__(self, directory, indices, train=True, fixed_seed=777): |
| self.directory, self.indices, self.train, self.fixed_seed = directory, list(indices), train, fixed_seed |
| self.rng = None |
| self.cache = {} |
|
|
| def __len__(self): |
| return len(self.indices) |
|
|
| def __getitem__(self, position): |
| import torch |
| index = self.indices[position] |
| if self.train: |
| if self.rng is None: |
| self.rng = np.random.default_rng(torch.initial_seed() % (2 ** 63)) |
| rng = self.rng |
| else: |
| rng = np.random.default_rng([self.fixed_seed, index]) |
| crop, target, kind = make_sample(rng, self.directory, index, self.cache) |
| rgb = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB) |
| return torch.from_numpy(np.ascontiguousarray(rgb.transpose(2, 0, 1))), torch.from_numpy(target), torch.tensor(kind) |
|
|
|
|
| def decode(heat_prob, offsets, threshold): |
| """3 x 3 local maxima above threshold -> [K, 2] crop px.""" |
| padded = np.pad(heat_prob, 1, constant_values=-1.0) |
| neighbours = np.max(np.stack([padded[dy:dy + CELLS, dx:dx + CELLS] for dy in range(3) for dx in range(3)]), axis=0) |
| ys, xs = np.nonzero((heat_prob >= neighbours) & (heat_prob >= threshold)) |
| return np.stack([(xs + offsets[ys, xs, 0]) * CELL - 0.5, (ys + offsets[ys, xs, 1]) * CELL - 0.5], axis=1) |
|
|
|
|
| def target_points(target): |
| ys, xs = np.nonzero(target[..., CH_POS] > 0.5) |
| return np.stack([(xs + target[ys, xs, CH_OFF_X]) * CELL - 0.5, (ys + target[ys, xs, CH_OFF_Y]) * CELL - 0.5], axis=1) |
|
|
|
|
| def match_count(predicted, truth, radius): |
| """Greedy one-to-one matching by distance -> number of matched pairs.""" |
| if len(predicted) == 0 or len(truth) == 0: |
| return 0 |
| distance = np.hypot(predicted[:, None, 0] - truth[None, :, 0], predicted[:, None, 1] - truth[None, :, 1]) |
| pi, ti = np.nonzero(distance <= radius) |
| used_p, used_t = set(), set() |
| for k in np.argsort(distance[pi, ti]): |
| if pi[k] not in used_p and ti[k] not in used_t: |
| used_p.add(pi[k]) |
| used_t.add(ti[k]) |
| return len(used_p) |
|
|
|
|
| def f1_scores(heat_prob, offsets, targets, kinds, thresholds, radius): |
| """Precision / recall / F1 on "red" crops per threshold and the best F1.""" |
| result, best = {}, -1.0 |
| for threshold in thresholds: |
| tp = fp = fn = 0 |
| for n in np.nonzero(kinds == KIND_CODES["red"])[0]: |
| predicted = decode(heat_prob[n], offsets[n], threshold) |
| truth = target_points(targets[n]) |
| m = match_count(predicted, truth, radius) |
| tp, fp, fn = tp + m, fp + len(predicted) - m, fn + len(truth) - m |
| precision, recall = tp / max(tp + fp, 1), tp / max(tp + fn, 1) |
| f1 = 2 * precision * recall / max(precision + recall, 1e-12) |
| result["t{:02d}".format(int(round(threshold * 100)))] = (precision, recall, f1) |
| best = max(best, f1) |
| result["best_f1"] = best |
| return result |
|
|
