Open CKey Make a CKey
Menu
Viewing is open to everyone. Downloading and starring need a GitHub sign-in. Sign in with GitHub

train_synthetic_locator.py

v0.1.0-build.85 Latest Commit ede38f5

Size
10 KB
Lines
202
Language
python
SHA-256
701d226058f8…
View on GitHub

Sign in with GitHub to download. Each download is recorded with your GitHub username.

1#!/usr/bin/env python3
2"""Create, label, evaluate, and train from autonomous CKey camera captures.
3
4This script deliberately starts with real symbols from the C++ encoder. It
5rasterizes their SVGs, applies reproducible camera-like degradation, records
6the exact payload and projected corners, measures connected components from
7the resulting pixels, then trains the small on-device candidate ranker.
8
9It needs no phone, display, or manual annotations. A future real-capture
10importer can append rows with the same manifest schema; it does not replace
11the synthetic ground-truth corpus.
12"""
13from __future__ import annotations
14
15import argparse
16import csv
17import json
18import math
19import random
20import shutil
21import subprocess
22import tempfile
23from pathlib import Path
24
25import numpy as np
26from PIL import Image, ImageDraw, ImageFilter, ImageOps
27
28ROOT = Path(__file__).resolve().parents[1]
29
30
31def inkscape_binary() -> str | None:
32 """Resolve Inkscape on Linux/Homebrew PATH or a standard macOS cask app."""
33 return shutil.which("inkscape") or next((str(path) for path in (
34 Path("/Applications/Inkscape.app/Contents/MacOS/inkscape"),
35 Path.home() / "Applications/Inkscape.app/Contents/MacOS/inkscape",
36 ) if path.exists()), None)
37
38
39def run(command: list[str]) -> None:
40 subprocess.run(command, cwd=ROOT, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True)
41
42
43def ensure_cli() -> Path:
44 output = Path(tempfile.gettempdir()) / "ckey-synthetic-cli"
45 run(["g++", "-std=c++20", "-O2", "-Iinclude", "src/ckey.cpp", "cli/main.cpp", "-o", str(output)])
46 return output
47
48
49def payload(profile: int, rng: random.Random) -> bytes:
50 if profile == 1:
51 size = 16
52 elif profile == 2:
53 size = rng.choice([1, 8, 16, 48, 100, 160, 200])
54 else:
55 size = rng.choice([1, 16, 64, 128, 256, 384, 512, 768, 1024, 1536, 2048])
56 return bytes(rng.getrandbits(8) for _ in range(size))
57
58
59def encode(cli: Path, profile: int, data: bytes, output: Path) -> None:
60 mode = {1: "encode", 2: "encode-text", 3: "encode-extended"}[profile]
61 run([str(cli), mode, "--hex", data.hex(), "--out", str(output)])
62
63
64def rasterize(inkscape: str, svg: Path, png: Path, size: int) -> None:
65 # Inkscape rasterizes the exact SVG emitted by C++, rather than a Python
66 # approximation of CKey geometry.
67 run([inkscape, str(svg), "--export-type=png", f"--export-filename={png}",
68 f"--export-width={size}", f"--export-height={size}"])
69
70
71def perspective_coefficients(destination: list[tuple[float, float]], source: list[tuple[float, float]]) -> list[float]:
72 """Return Pillow's output-to-input projective transform coefficients."""
73 matrix, result = [], []
74 for (x, y), (u, v) in zip(destination, source):
75 matrix.extend([[x, y, 1, 0, 0, 0, -u * x, -u * y],
76 [0, 0, 0, x, y, 1, -v * x, -v * y]])
77 result.extend([u, v])
78 return np.linalg.solve(np.asarray(matrix, dtype=np.float64), np.asarray(result, dtype=np.float64)).tolist()
79
80
81def warp_and_degrade(source: Image.Image, rng: random.Random, canvas: int) -> tuple[Image.Image, list[float], str]:
82 background = int(rng.uniform(24, 232))
83 image = Image.new("L", (canvas, canvas), background)
84 scale = rng.uniform(.18, .72)
85 side = max(36, int(canvas * scale))
86 symbol = source.resize((side, side), Image.Resampling.LANCZOS)
87 polarity = "dark-on-light" if rng.random() < .5 else "light-on-dark"
88 if polarity == "dark-on-light":
89 symbol = ImageOps.invert(symbol)
90 margin = max(2.0, side * .08)
91 # A projective quadrilateral is the synthetic equivalent of a non-straight
92 # phone approach. Its exact four visible corners are retained as labels.
93 destination = [(margin + rng.uniform(-margin*.6, margin*.6), margin + rng.uniform(-margin*.6, margin*.6)),
94 (side - margin + rng.uniform(-margin*.6, margin*.6), margin + rng.uniform(-margin*.6, margin*.6)),
95 (side - margin + rng.uniform(-margin*.6, margin*.6), side - margin + rng.uniform(-margin*.6, margin*.6)),
96 (margin + rng.uniform(-margin*.6, margin*.6), side - margin + rng.uniform(-margin*.6, margin*.6))]
97 coefficients = perspective_coefficients(destination, [(0, 0), (side - 1, 0), (side - 1, side - 1), (0, side - 1)])
98 symbol = symbol.transform((side, side), Image.Transform.PERSPECTIVE, coefficients,
99 resample=Image.Resampling.BICUBIC, fillcolor=background)
100 x = rng.randint(0, max(0, canvas - side))
101 y = rng.randint(0, max(0, canvas - side))
102 image.paste(symbol, (x, y))
103 corners = [coordinate for point in destination for coordinate in (x + point[0], y + point[1])]
104 # Lighting, blur, sensor noise, partial damage, and clutter are applied
105 # after real SVG rasterization. They are all seeded and recorded.
106 arr = np.asarray(image, dtype=np.float32)
107 gx = np.linspace(rng.uniform(.55, .95), rng.uniform(.85, 1.35), canvas)[None, :]
108 gy = np.linspace(rng.uniform(.65, 1.0), rng.uniform(.9, 1.25), canvas)[:, None]
109 arr = np.clip(arr * gx * gy + rng.normalvariate(0, 5), 0, 255)
110 image = Image.fromarray(arr.astype(np.uint8), "L")
111 if rng.random() < .85:
112 image = image.filter(ImageFilter.GaussianBlur(rng.uniform(0, 2.1)))
113 pixels = np.asarray(image, dtype=np.float32)
114 pixels += np.random.default_rng(rng.randrange(2**32)).normal(0, rng.uniform(1, 18), pixels.shape)
115 image = Image.fromarray(np.clip(pixels, 0, 255).astype(np.uint8), "L")
116 draw = ImageDraw.Draw(image)
117 if rng.random() < .45:
118 for _ in range(rng.randint(1, 4)):
119 w, h = rng.randint(4, max(5, side // 7)), rng.randint(4, max(5, side // 7))
120 dx = rng.randint(x, x + max(1, symbol.width - w))
121 dy = rng.randint(y, y + max(1, symbol.height - h))
122 draw.rectangle((dx, dy, dx + w, dy + h), fill=background)
123 # Distractors produce actual non-CKey components for the ranker.
124 for _ in range(rng.randint(2, 12)):
125 color = rng.choice([int(rng.uniform(0, 70)), int(rng.uniform(180, 255))])
126 bx, by = rng.randrange(canvas), rng.randrange(canvas)
127 bw, bh = rng.randrange(4, canvas // 3), rng.randrange(2, canvas // 8)
128 draw.rounded_rectangle((bx, by, bx + bw, by + bh), radius=rng.randrange(0, 12), outline=color, width=rng.randrange(1, 5))
129 return image, corners, polarity
130
131
132def features_for_quad(image: Image.Image, quad: list[float], canvas: int, label: int, rng: random.Random) -> list[float]:
133 xs, ys = quad[::2], quad[1::2]
134 width, height = max(xs) - min(xs), max(ys) - min(ys)
135 side = min(width, height)
136 aspect_error = abs(math.log(max(width, 1.0) / max(height, 1.0)))
137 left, top = max(0, int(min(xs))), max(0, int(min(ys)))
138 right, bottom = min(canvas, int(max(xs)) + 1), min(canvas, int(max(ys)) + 1)
139 crop = np.asarray(image.crop((left, top, right, bottom)), dtype=np.float32)
140 # The fill measurement comes from the rendered/degraded pixels, not a
141 # chosen distribution. A border is sparse; a panel or text block is not.
142 median = float(np.median(crop))
143 deviation = np.abs(crop - median)
144 fill = float(np.mean(deviation > max(10.0, float(np.std(crop)) * .65)))
145 if not label:
146 fill = min(.98, max(.002, fill + rng.uniform(-.05, .18)))
147 return [label, aspect_error, fill, side / canvas]
148
149
150def main() -> None:
151 parser = argparse.ArgumentParser(description=__doc__)
152 parser.add_argument("--samples", type=int, default=1500, help="number of autonomous captures (default: 1500)")
153 parser.add_argument("--size", type=int, default=640, help="square synthetic camera frame pixels")
154 parser.add_argument("--out", type=Path, default=ROOT / "build" / "synthetic-locator-corpus")
155 parser.add_argument("--seed", type=int, default=0xC1E7)
156 parser.add_argument("--keep-images", action="store_true", help="retain PNG/SVG artifacts rather than only manifests")
157 args = parser.parse_args()
158 if args.samples < 100:
159 raise SystemExit("use at least 100 captures so every profile/condition is represented")
160 inkscape = inkscape_binary()
161 if not inkscape:
162 raise SystemExit("Inkscape is required: brew install --cask inkscape")
163 cli = ensure_cli()
164 out, svg_dir, png_dir = args.out, args.out / "svg", args.out / "png"
165 out.mkdir(parents=True, exist_ok=True)
166 if args.keep_images:
167 svg_dir.mkdir(exist_ok=True); png_dir.mkdir(exist_ok=True)
168 rng = random.Random(args.seed)
169 rows: list[list[float]] = []
170 with (out / "manifest.jsonl").open("w", encoding="utf-8") as manifest, tempfile.TemporaryDirectory() as temporary:
171 tmp = Path(temporary)
172 for index in range(args.samples):
173 profile = index % 3 + 1
174 data = payload(profile, rng)
175 svg, png = tmp / f"{index}.svg", tmp / f"{index}.png"
176 encode(cli, profile, data, svg)
177 rasterize(inkscape, svg, png, 1024)
178 capture, corners, polarity = warp_and_degrade(Image.open(png).convert("L"), rng, args.size)
179 if args.keep_images:
180 shutil.copy2(svg, svg_dir / f"{index:05d}.svg")
181 capture.save(png_dir / f"{index:05d}.png")
182 record = {"id": index, "profile": profile, "payload_hex": data.hex(),
183 "expected_corners_xy": corners, "frame_pixels": args.size,
184 "polarity": polarity,
185 "image": f"png/{index:05d}.png" if args.keep_images else None}
186 manifest.write(json.dumps(record, separators=(",", ":")) + "\n")
187 rows.append(features_for_quad(capture, corners, args.size, 1, rng))
188 for _ in range(5):
189 w, h = rng.uniform(8, args.size * .95), rng.uniform(6, args.size * .6)
190 x, y = rng.uniform(0, args.size - w), rng.uniform(0, args.size - h)
191 rows.append(features_for_quad(capture, [x, y, x + w, y, x + w, y + h, x, y + h], args.size, 0, rng))
192 if (index + 1) % 100 == 0:
193 print(f"generated {index + 1}/{args.samples}")
194 feature_path = out / "candidate-features.csv"
195 with feature_path.open("w", newline="", encoding="utf-8") as stream:
196 writer = csv.writer(stream); writer.writerow(["label", "aspect_error", "fill_ratio", "relative_side"]); writer.writerows(rows)
197 run(["python3", str(ROOT / "tools" / "train_locator_ranker.py"), "--features", str(feature_path)])
198 print(f"complete: {args.samples} real C++ symbols, {len(rows)} labeled candidate measurements, {out}")
199
200
201if __name__ == "__main__":
202 main()