8f43927575
examples/ is based on KasperskyLab/knp at 55470ed6 with local changes to the mnist-learn altai/blifat models, cuda-backend CMake and model_loader. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
66 lines
3.2 KiB
Python
66 lines
3.2 KiB
Python
"""Prepare reproducible CIFAR-10 cat/dog splits for the KNP example (numpy only)."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import tarfile
|
|
import urllib.request
|
|
|
|
import numpy as np
|
|
|
|
URL = 'https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz'
|
|
MD5 = 'c32a1d4ab5d03f1284b67883e8d87530'
|
|
|
|
|
|
def encode(rgb):
|
|
# Fixed, untrained preprocessing. Separate brightness and darkness channels.
|
|
gray = ((rgb[:, 0].astype(np.uint32) * 299 + rgb[:, 1].astype(np.uint32) * 587
|
|
+ rgb[:, 2].astype(np.uint32) * 114 + 500) // 1000).astype(np.uint8)
|
|
return np.concatenate([gray, 255 - gray], axis=1)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('output', type=Path)
|
|
a = ap.parse_args()
|
|
a.output.mkdir(parents=True, exist_ok=True)
|
|
archive = a.output / 'cifar-10-binary.tar.gz'
|
|
if not archive.exists():
|
|
urllib.request.urlretrieve(URL, archive)
|
|
assert hashlib.md5(archive.read_bytes()).hexdigest() == MD5, 'Archive checksum mismatch'
|
|
with tarfile.open(archive) as tar:
|
|
def read(names):
|
|
data = np.concatenate([np.frombuffer(tar.extractfile('cifar-10-batches-bin/' + name).read(),
|
|
dtype=np.uint8).reshape(-1, 3073) for name in names])
|
|
ids = np.flatnonzero(np.isin(data[:, 0], [3, 5]))
|
|
return data[ids, 1:].reshape(-1, 3, 1024), (data[ids, 0] == 5).astype(np.uint8), ids
|
|
rgb, labels, ids = read([f'data_batch_{i}.bin' for i in range(1, 6)])
|
|
test_rgb, test_labels, test_ids = read(['test_batch.bin'])
|
|
rng = np.random.default_rng(20261014)
|
|
train, val = [], []
|
|
for label in (0, 1):
|
|
indices = rng.permutation(np.flatnonzero(labels == label))
|
|
train.extend(indices[:4000])
|
|
val.extend(indices[4000:])
|
|
train = rng.permutation(train)
|
|
val = rng.permutation(val)
|
|
test = np.random.default_rng(20261015).permutation(len(test_labels))
|
|
for name, x, y, source in [('train', rgb[train], labels[train], ids[train]),
|
|
('val', rgb[val], labels[val], ids[val]),
|
|
('test', test_rgb[test], test_labels[test], test_ids[test])]:
|
|
encode(x).tofile(a.output / f'{name}.bin')
|
|
np.savetxt(a.output / f'{name}.target', y, fmt='%d')
|
|
np.savez_compressed(a.output / f'{name}-original.npz', rgb=x, labels=y, source_indices=source)
|
|
for name, indices, inference in [('smoke', train[:100], val[:100]), ('validation', train, val)]:
|
|
encode(np.concatenate([rgb[indices], rgb[inference]])).tofile(a.output / f'{name}.bin')
|
|
np.savetxt(a.output / f'{name}.target', np.concatenate([labels[indices], labels[inference]]), fmt='%d')
|
|
manifest = {'source': URL, 'md5': MD5, 'labels': {'0': 'cat', '1': 'dog'},
|
|
'seed': 20261014, 'train': 8000, 'validation': 2000, 'test': 2000,
|
|
'encoding': '32x32 grayscale integer BT.601; concatenate intensity and 255-intensity',
|
|
'features': 2048, 'train_validation_disjoint': bool(not set(train) & set(val))}
|
|
(a.output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
|
|
print(json.dumps(manifest), flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|