File size: 1,281 Bytes
fd684da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | """Convert the official CIFAR-10 training split into the wARTok PNG layout."""
import argparse
import json
from pathlib import Path
import numpy as np
from PIL import Image
from torchvision.datasets import CIFAR10
def main():
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--download-dir',default='./data/downloads')
p.add_argument('--output',default='./data/cifar10')
p.add_argument('--download',action='store_true')
p.add_argument('--reference-images',help='Optional NPZ of the same 50K training images for ADM statistics')
a=p.parse_args(); out=Path(a.output)
if out.exists() and any(out.iterdir()): p.error('output must be empty to prevent mixing splits')
ds=CIFAR10(a.download_dir,train=True,download=a.download)
assert len(ds)==50000
labels=[]
for i,(image,label) in enumerate(ds):
name=f'{i//1000:05d}/img{i:08d}.png'
path=out/name; path.parent.mkdir(parents=True,exist_ok=True)
image.save(path); labels.append([name,label])
(out/'dataset.json').write_text(json.dumps({'labels':labels}))
if a.reference_images:
with open(a.reference_images,'xb') as f: np.savez_compressed(f,ds.data)
print(f'Wrote {len(ds)} training images to {out}')
if __name__=='__main__': main()
|