65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
"""Build ON/OFF-coded MNIST datasets for mnist_learn_example (KNP_MNIST_INPUT_SIZE=1568).
|
|
|
|
Each image becomes 1568 bytes: 784 ON pixels (unchanged) followed by 784 OFF pixels.
|
|
OFF modes:
|
|
full - 255 - p everywhere;
|
|
half - (255 - p) / 2 everywhere;
|
|
local - 255 - p only within `--radius` pixels of the digit (p > 0), zero elsewhere.
|
|
|
|
usage: python prepare_onoff.py <src.bin> <src.target> <out_prefix> --mode local [--radius 2]
|
|
Writes <out_prefix>.bin and <out_prefix>.target (labels copied as is).
|
|
"""
|
|
import argparse
|
|
import shutil
|
|
|
|
import numpy as np
|
|
|
|
SIDE = 28
|
|
|
|
|
|
def dilate(mask, radius):
|
|
out = mask.copy()
|
|
for dy in range(-radius, radius + 1):
|
|
for dx in range(-radius, radius + 1):
|
|
if dy * dy + dx * dx > radius * radius:
|
|
continue
|
|
shifted = np.zeros_like(mask)
|
|
ys, yd = (slice(dy, None), slice(None, -dy or None)) if dy >= 0 else (slice(None, dy), slice(-dy, None))
|
|
xs, xd = (slice(dx, None), slice(None, -dx or None)) if dx >= 0 else (slice(None, dx), slice(-dx, None))
|
|
shifted[:, ys, xs] = mask[:, yd, xd]
|
|
out |= shifted
|
|
return out
|
|
|
|
|
|
def off_channel(imgs, mode, radius):
|
|
inv = 255 - imgs.astype(np.int32)
|
|
if mode == "full":
|
|
off = inv
|
|
elif mode == "half":
|
|
off = inv // 2
|
|
elif mode == "local":
|
|
off = inv * dilate(imgs > 0, radius)
|
|
else:
|
|
raise ValueError(mode)
|
|
return off.astype(np.uint8)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("images")
|
|
ap.add_argument("labels")
|
|
ap.add_argument("out")
|
|
ap.add_argument("--mode", choices=["full", "half", "local"], required=True)
|
|
ap.add_argument("--radius", type=int, default=2)
|
|
a = ap.parse_args()
|
|
imgs = np.fromfile(a.images, dtype=np.uint8).reshape(-1, SIDE, SIDE)
|
|
off = off_channel(imgs, a.mode, a.radius)
|
|
out = np.concatenate([imgs.reshape(len(imgs), -1), off.reshape(len(imgs), -1)], axis=1)
|
|
out.tofile(a.out + ".bin")
|
|
shutil.copyfile(a.labels, a.out + ".target")
|
|
on_spk, off_spk = imgs.mean() / 255 * 10 * SIDE * SIDE, off.mean() / 255 * 10 * SIDE * SIDE
|
|
print(f"{a.out}: {len(out)} images, ~{on_spk:.0f} ON / ~{off_spk:.0f} OFF spikes per image")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|