62 lines
2.8 KiB
Python
62 lines
2.8 KiB
Python
"""CIFAR-10 vehicles (airplane, automobile, ship, truck → 0..3) from the Hugging Face parquet copy (uoft-cs/cifar10).
|
|
|
|
train: first 4000 training images of each class (16000), val: the last 1000 of each class (4000),
|
|
test: the official test images of these classes (4000) — kept for one final evaluation.
|
|
usage: ~/anvenv/bin/python prepare_vehicles.py → /home/int/run/cifar/vehicles/{train,val,test}.npz
|
|
"""
|
|
import io
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow.parquet as pq
|
|
from PIL import Image
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('--source', type=Path, default=Path.home() / 'run/cifar')
|
|
ap.add_argument('--out', type=Path, default=Path.home() / 'run/cifar/vehicles')
|
|
a = ap.parse_args()
|
|
src, out = a.source, a.out
|
|
out.mkdir(parents=True, exist_ok=True)
|
|
CLASSES = {0: 0, 1: 1, 8: 2, 9: 3} # CIFAR label → our label
|
|
NAMES = ['airplane', 'automobile', 'ship', 'truck']
|
|
|
|
|
|
def load(name):
|
|
t = pq.read_table(src / f'hf_{name}.parquet').to_pydict()
|
|
imgs, labels = [], []
|
|
for img, lab in zip(t['img'], t['label']):
|
|
if lab in CLASSES:
|
|
imgs.append(np.asarray(Image.open(io.BytesIO(img['bytes'])).convert('RGB')).transpose(2, 0, 1))
|
|
labels.append(CLASSES[lab])
|
|
return np.stack(imgs), np.array(labels, np.uint8)
|
|
|
|
|
|
x, y = load('train')
|
|
idx = {c: np.flatnonzero(y == c) for c in range(4)}
|
|
train = np.concatenate([idx[c][:4000] for c in range(4)])
|
|
val = np.concatenate([idx[c][4000:] for c in range(4)])
|
|
rng = np.random.default_rng(20261015)
|
|
train, val = rng.permutation(train), rng.permutation(val)
|
|
np.savez_compressed(out / 'train.npz', rgb=x[train], labels=y[train])
|
|
np.savez_compressed(out / 'val.npz', rgb=x[val], labels=y[val])
|
|
xt, yt = load('test')
|
|
p = rng.permutation(len(yt))
|
|
np.savez_compressed(out / 'test.npz', rgb=xt[p], labels=yt[p])
|
|
print('train', len(train), np.bincount(y[train]), 'val', len(val), 'test', len(yt), np.bincount(yt), NAMES)
|
|
|
|
# Native input format used by altai_infer; mirrors only augment the training split.
|
|
for split in ('train', 'val', 'test'):
|
|
with np.load(out / f'{split}.npz') as saved:
|
|
images, labels = saved['rgb'], saved['labels']
|
|
images.tofile(out / f'{split}_images.bin')
|
|
(out / f'{split}_labels.txt').write_text('\n'.join(map(str, labels.tolist())) + '\n')
|
|
if split == 'train':
|
|
images[:, :, :, ::-1].copy().tofile(out / 'train_flip_images.bin')
|
|
manifest = {'source': 'uoft-cs/cifar10 parquet', 'seed': 20261015, 'classes': NAMES,
|
|
'split': 'first 4000 per class train, last 1000 validation, official test reserved',
|
|
'files': {p.name: {'bytes': p.stat().st_size, 'sha256': hashlib.sha256(p.read_bytes()).hexdigest()}
|
|
for p in sorted(out.iterdir()) if p.suffix in {'.bin', '.txt', '.npz'}}}
|
|
(out / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
|