AWWL-Diff/data.py

# prepare_ref_data.py
import os
from torchvision.datasets import CIFAR10
from tqdm import tqdm
from PIL import Image


def save_cifar10_images(output_dir="cifar10_train_ref"):
    # Tải dataset CIFAR-10 (Train split)
    dataset = CIFAR10(root="./data", train=True, download=True)

    os.makedirs(output_dir, exist_ok=True)
    print(f"Saving {len(dataset)} CIFAR-10 training images to '{output_dir}'...")

    for idx, (img, label) in enumerate(tqdm(dataset)):
        # img là PIL Image
        img.save(f"{output_dir}/{idx:05d}.png")

    print("✅ Done! Reference images are ready.")


if __name__ == "__main__":
    save_cifar10_images()