82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
"""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()
|