| """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() |
|
|
