# save_examples.py
import os, torch
from torchvision.datasets import CIFAR100
from torchvision.utils import save_image
from src.dataset import get_cifar100_transforms, MEAN, STD

os.makedirs("assets/examples", exist_ok=True)

# Load CIFAR100 test set
dataset = CIFAR100(root="data", train=False, download=True,
                   transform=get_cifar100_transforms(augment=False))

for i in range(15):  # pick first 15 test images
    img, label = dataset[i]
    img_vis = img * torch.tensor(STD).view(3,1,1) + torch.tensor(MEAN).view(3,1,1)
    path = f"assets/examples/{dataset.classes[label]}.png"
    save_image(img_vis, path)

print("✅ Saved example images to assets/examples/")