File size: 3,844 Bytes
3ce19a2 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Download CIFAR-10 and extract images for training + FID evaluation.
Usage:
python prepare_cifar10.py [--data_root ./cifar10_data]
Creates:
<data_root>/cifar-10-batches-py/ (pickle files used by data.py)
<data_root>/img/ (PNG images for FID computation)
"""
import argparse
import os
import pickle
import shutil
import tarfile
import urllib.request
import numpy as np
from PIL import Image
CIFAR10_URL = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
def is_valid_tar_gz(path: str) -> bool:
if not os.path.isfile(path):
return False
try:
with tarfile.open(path, "r:gz") as tf:
return tf.getmember("cifar-10-batches-py") is not None
except Exception:
return False
def ensure_clean_tarball(path: str):
if os.path.isfile(path) and not is_valid_tar_gz(path):
print(f"Corrupted archive detected, removing: {path}")
os.remove(path)
def has_complete_batches(batch_dir: str) -> bool:
required = [
"data_batch_1",
"data_batch_2",
"data_batch_3",
"data_batch_4",
"data_batch_5",
"test_batch",
"batches.meta",
]
return all(os.path.isfile(os.path.join(batch_dir, name)) for name in required)
def download_cifar10(data_root: str):
tar_path = os.path.join(data_root, "cifar-10-python.tar.gz")
batch_dir = os.path.join(data_root, "cifar-10-batches-py")
if os.path.isdir(batch_dir) and has_complete_batches(batch_dir):
print(f"Already exists: {batch_dir}")
return
if os.path.isdir(batch_dir):
print(f"Incomplete batch directory detected, removing: {batch_dir}")
shutil.rmtree(batch_dir, ignore_errors=True)
os.makedirs(data_root, exist_ok=True)
ensure_clean_tarball(tar_path)
if not os.path.isfile(tar_path):
print(f"Downloading CIFAR-10 to {tar_path} ...")
urllib.request.urlretrieve(CIFAR10_URL, tar_path)
print("Download complete.")
ensure_clean_tarball(tar_path)
print("Extracting ...")
try:
with tarfile.open(tar_path, "r:gz") as tf:
tf.extractall(data_root)
except Exception as e:
print(f"Extraction failed ({type(e).__name__}); re-downloading archive once ...")
if os.path.isdir(batch_dir):
shutil.rmtree(batch_dir, ignore_errors=True)
if os.path.isfile(tar_path):
os.remove(tar_path)
urllib.request.urlretrieve(CIFAR10_URL, tar_path)
with tarfile.open(tar_path, "r:gz") as tf:
tf.extractall(data_root)
print(f"Extracted to {batch_dir}")
def extract_images(data_root: str):
"""Extract all training images as PNGs into <data_root>/img/ for FID."""
img_dir = os.path.join(data_root, "img")
if os.path.isdir(img_dir) and len(os.listdir(img_dir)) >= 45000:
print(f"Image folder already populated: {img_dir}")
return
os.makedirs(img_dir, exist_ok=True)
batch_dir = os.path.join(data_root, "cifar-10-batches-py")
idx = 0
for batch_id in range(1, 6):
path = os.path.join(batch_dir, f"data_batch_{batch_id}")
with open(path, "rb") as f:
batch = pickle.load(f, encoding="bytes")
images = batch[b"data"].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
for img_np in images:
Image.fromarray(img_np).save(
os.path.join(img_dir, f"{idx:06d}.png"))
idx += 1
print(f"Extracted {idx} training images to {img_dir}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_root", type=str, default="./cifar10_data")
args = parser.parse_args()
download_cifar10(args.data_root)
extract_images(args.data_root)
print("Done. Use --data_root", args.data_root, "when launching training.")
|