41 lines
1.7 KiB
Python
41 lines
1.7 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
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow.parquet as pq
|
|
from PIL import Image
|
|
|
|
src, out = Path('/home/int/run/cifar'), Path('/home/int/run/cifar/vehicles')
|
|
out.mkdir(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)
|