diff --git a/augment/augment_mnist.py b/augment/augment_mnist.py new file mode 100644 index 0000000..5b657c2 --- /dev/null +++ b/augment/augment_mnist.py @@ -0,0 +1,82 @@ +"""Training-set augmentation for mnist_learn_example datasets. + +Input file layout (as in ~/run/exp): , 784 bytes per image. +Output: each of the first `--base` training images is kept and followed by `--copies` random +variants (shift up to ±2 px, rotation up to ±12°, scale 0.9–1.1, optional 1 px thicker/thinner +stroke); the training part is shuffled; the eval part is copied unchanged. + +usage: python augment_mnist.py --train 30000 --base 30000 --copies 2 +Writes .bin/.target, where N = base * (copies + 1) is the new training size, +so that run.sh can use DATA= and n=N. +""" +import argparse + +import numpy as np +from PIL import Image + +SIDE = 28 +CROSS = [(0, 1), (0, -1), (1, 0), (-1, 0)] + + +def shift(a, dy, dx): + return np.roll(np.roll(a, dy, 0), dx, 1) + + +def augment(img, rng): + angle = rng.uniform(-12, 12) + scale = rng.uniform(0.9, 1.1) + dy, dx = rng.integers(-2, 3, size=2) + out = Image.fromarray(img).rotate(angle, resample=Image.BILINEAR, center=(13.5, 13.5)) + if scale != 1: + size = max(1, round(SIDE * scale)) + scaled = out.resize((size, size), Image.BILINEAR) + canvas = Image.new("L", (SIDE, SIDE)) + off = (SIDE - size) // 2 + canvas.paste(scaled, (off, off)) if size <= SIDE else canvas.paste(scaled.crop((-off, -off, -off + SIDE, -off + SIDE))) + out = canvas + a = np.asarray(out, dtype=np.uint8) + stroke = rng.choice(["same", "thicker", "thinner"], p=[0.5, 0.25, 0.25]) + if stroke == "thicker": + a = np.maximum.reduce([a] + [shift(a, y, x) for y, x in CROSS]) + elif stroke == "thinner": + a = np.minimum.reduce([a] + [shift(a, y, x) for y, x in CROSS]) + # Shift without wrap-around. + res = np.zeros_like(a) + ys, yd = (slice(dy, None), slice(None, SIDE - dy)) if dy >= 0 else (slice(None, dy), slice(-dy, None)) + xs, xd = (slice(dx, None), slice(None, SIDE - dx)) if dx >= 0 else (slice(None, dx), slice(-dx, None)) + res[ys, xs] = a[yd, xd] + return res + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("images") + ap.add_argument("labels") + ap.add_argument("out") + ap.add_argument("--train", type=int, required=True, help="training images in the source file") + ap.add_argument("--base", type=int, required=True, help="how many source training images to use") + ap.add_argument("--copies", type=int, default=2) + ap.add_argument("--seed", type=int, default=20261014) + a = ap.parse_args() + rng = np.random.default_rng(a.seed) + imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE) + labels = open(a.labels).read().split() + train_imgs, train_labels = [], [] + for i in range(a.base): + train_imgs.append(imgs[i]) + train_labels.append(labels[i]) + for _ in range(a.copies): + train_imgs.append(augment(imgs[i], rng)) + train_labels.append(labels[i]) + order = rng.permutation(len(train_imgs)) + out_imgs = np.concatenate([np.stack(train_imgs)[order], imgs[a.train:]]) + out_labels = [train_labels[j] for j in order] + labels[a.train:] + n = len(train_imgs) + out_imgs.tofile(f"{a.out}{n}.bin") + with open(f"{a.out}{n}.target", "w") as f: + f.write("\n".join(out_labels) + "\n") + print(f"{a.out}{n}: {n} training + {len(imgs) - a.train} eval images") + + +if __name__ == "__main__": + main() diff --git a/augment/tta_sim.py b/augment/tta_sim.py new file mode 100644 index 0000000..30ac47f --- /dev/null +++ b/augment/tta_sim.py @@ -0,0 +1,65 @@ +"""Test-time shift voting in the offline AltAI simulator. + +Each image is fed several times with small shifts; votes of all copies are summed. +Also measures robustness: the test images themselves are shifted / thickened first. + +usage: python tta_sim.py --first 30000 --test 2000 +""" +import argparse +import sys + +import numpy as np + +sys.path.insert(0, "/mnt/c/knp/viz") +import altai_sim as s # noqa: E402 + +CROSS = [(0, 0), (0, 1), (0, -1), (1, 0), (-1, 0)] +SQUARE = [(dy, dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1)] + + +def shift(a, dy, dx): + return np.roll(np.roll(a, dy, -2), dx, -1) + + +def grow(a): + return np.maximum.reduce([shift(a, dy, dx) for dy, dx in CROSS]) + + +def thin(a): + return np.minimum.reduce([shift(a, dy, dx) for dy, dx in CROSS]) + + +def evaluate(w, thr, imgs, labels, offsets, inhib): + sim = s.Sim(w, thr, inhib) + correct = silent = 0 + for img, lab in zip(imgs, labels): + votes = sum(sim.run_image(shift(img, dy, dx).reshape(-1))[0] for dy, dx in offsets) + p = s.predict(votes) + correct += p == lab + silent += p < 0 + return 100 * correct / len(labels), 100 * silent / len(labels) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("dump") + ap.add_argument("images") + ap.add_argument("labels") + ap.add_argument("--first", type=int, default=30000) + ap.add_argument("--test", type=int, default=2000) + ap.add_argument("--inhib", type=float, default=255) + a = ap.parse_args() + w, thr = s.load_dump(a.dump) + imgs, labels = s.load_images(a.images, a.labels, a.first, a.test) + imgs = imgs.reshape(-1, 28, 28).astype(np.int32) + variants = {"as is": imgs, "shift 2px": shift(imgs, 0, 2), "thicker 1px": grow(imgs), "thinner 1px": thin(imgs)} + for vote_name, offsets in [("1 copy", [(0, 0)]), ("5 copies (cross)", CROSS), ("9 copies (3x3)", SQUARE)]: + row = [] + for name, x in variants.items(): + acc, sil = evaluate(w, thr, x.clip(0, 255).astype(np.uint8), labels, offsets, a.inhib) + row.append(f"{name} {acc:5.1f}%") + print(f"{vote_name:17s} | " + " | ".join(row), flush=True) + + +if __name__ == "__main__": + main()