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>
106 lines
4.8 KiB
Python
106 lines
4.8 KiB
Python
"""Small pilot using JPEG previews of uoft-cs/cifar10, not canonical CIFAR bytes."""
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import hashlib
|
|
from io import BytesIO
|
|
import json
|
|
from pathlib import Path
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
from prepare_data import encode
|
|
|
|
root = Path('/home/int/run/cats-dogs/data')
|
|
cache = root / 'preview-cache'
|
|
cache.mkdir(exist_ok=True)
|
|
|
|
|
|
def request(url):
|
|
for attempt in range(8):
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=45) as response:
|
|
return response.read()
|
|
except Exception as error:
|
|
if attempt == 7:
|
|
raise
|
|
delay = min(45, 3 * (attempt + 1))
|
|
if isinstance(error, urllib.error.HTTPError) and error.code == 429:
|
|
delay = max(delay, int(error.headers.get('Retry-After', '15')))
|
|
time.sleep(delay)
|
|
|
|
|
|
def rows(args):
|
|
split, offset = args
|
|
cached = cache / f'rows-{split}-{offset}.json'
|
|
if cached.exists():
|
|
return split, json.loads(cached.read_text())['rows']
|
|
params = urllib.parse.urlencode(dict(dataset='uoft-cs/cifar10', config='plain_text',
|
|
split=split, offset=offset, length=100))
|
|
response = request('https://datasets-server.huggingface.co/rows?' + params)
|
|
cached.write_bytes(response)
|
|
time.sleep(0.5)
|
|
return split, json.loads(response)['rows']
|
|
|
|
|
|
catalog = {'train': [], 'test': []}
|
|
tasks = [('train', i) for i in range(0, 6000, 100)] + [('test', i) for i in range(0, 3000, 100)]
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
for split, batch in pool.map(rows, tasks):
|
|
catalog[split].extend(r for r in batch if r['row']['label'] in (3, 5))
|
|
selected = {'train': [], 'val': [], 'test': []}
|
|
rng = np.random.default_rng(20261014)
|
|
for label in (3, 5):
|
|
candidates = sorted([r for r in catalog['train'] if r['row']['label'] == label], key=lambda r: r['row_idx'])
|
|
candidates = [candidates[i] for i in rng.permutation(len(candidates))]
|
|
assert len(candidates) >= 500
|
|
selected['train'].extend(candidates[:400])
|
|
selected['val'].extend(candidates[400:500])
|
|
candidates = sorted([r for r in catalog['test'] if r['row']['label'] == label], key=lambda r: r['row_idx'])
|
|
assert len(candidates) >= 200
|
|
selected['test'].extend(candidates[:200])
|
|
for split in selected:
|
|
selected[split] = [selected[split][i] for i in rng.permutation(len(selected[split]))]
|
|
|
|
|
|
def image(args):
|
|
split, row = args
|
|
path = cache / f'{split}-{row["row_idx"]}.jpg'
|
|
if not path.exists():
|
|
path.write_bytes(request(row['row']['img']['src']))
|
|
img = np.asarray(Image.open(BytesIO(path.read_bytes())).convert('RGB'))
|
|
assert img.shape == (32, 32, 3)
|
|
return img.transpose(2, 0, 1).reshape(3, 1024)
|
|
|
|
|
|
arrays, labels = {}, {}
|
|
provenance = {}
|
|
with ThreadPoolExecutor(max_workers=4) as pool:
|
|
for split in selected:
|
|
original_split = 'test' if split == 'test' else 'train'
|
|
arrays[split] = np.stack(list(pool.map(image, [(original_split, row) for row in selected[split]])))
|
|
labels[split] = np.array([int(row['row']['label'] == 5) for row in selected[split]], dtype=np.uint8)
|
|
provenance[split] = [row['row_idx'] for row in selected[split]]
|
|
encode(arrays[split]).tofile(root / f'{split}.bin')
|
|
np.savetxt(root / f'{split}.target', labels[split], fmt='%d')
|
|
np.savez_compressed(root / f'{split}-original.npz', rgb=arrays[split], labels=labels[split],
|
|
source_indices=provenance[split])
|
|
print(split, len(labels[split]), 'ready', flush=True)
|
|
for name, first, second in [('smoke', ('train', 100), ('val', 100)),
|
|
('validation', ('train', 800), ('val', 200)),
|
|
('train_test', ('train', 800), ('test', 400))]:
|
|
encode(np.concatenate([arrays[first[0]][:first[1]], arrays[second[0]][:second[1]]])).tofile(root / f'{name}.bin')
|
|
np.savetxt(root / f'{name}.target', np.concatenate([labels[first[0]][:first[1]], labels[second[0]][:second[1]]]), fmt='%d')
|
|
assert not set(provenance['train']) & set(provenance['val'])
|
|
manifest = {'source': 'https://huggingface.co/datasets/uoft-cs/cifar10',
|
|
'variant': 'JPEG previews from datasets-server API; not original binary archive',
|
|
'labels': {'0': 'cat', '1': 'dog'}, 'seed': 20261014,
|
|
'train': 800, 'validation': 200, 'test': 400, 'features': 2048,
|
|
'encoding': 'integer grayscale BT.601, concatenate brightness and darkness',
|
|
'source_indices': provenance, 'train_validation_disjoint': True,
|
|
'encoded_sha256': {s: hashlib.sha256((root / f'{s}.bin').read_bytes()).hexdigest() for s in selected}}
|
|
(root / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
|
|
print('Pilot dataset ready', flush=True)
|