Augmentation experiment: training-set augmenter and test-time shift voting for MNIST

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

augment/augment_mnist.py

@@ -0,0 +1,82 @@
+"""Training-set augmentation for mnist_learn_example datasets.
+
+Input file layout (as in ~/run/exp): <train images><eval images>, 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 <src.bin> <src.target> <out_prefix> --train 30000 --base 30000 --copies 2
+Writes <out_prefix><N>.bin/.target, where N = base * (copies + 1) is the new training size,
+so that run.sh can use DATA=<out_prefix> 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()

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 <dump_dir> <images.bin> <labels.target> --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()