Update app.py

app.py

@@ -73,24 +73,26 @@ def show_results(input_image: Image.Image):
 
     return preds, perf_plot, acc_plot, test_acc_text
 
-# === Prepare CIFAR-10 Sample Gallery ===
-# Download CIFAR-10 test set and save a few sample images
+# === Prepare CIFAR-10 Sample Gallery (one per class) ===
 sample_dir = "examples"
 os.makedirs(sample_dir, exist_ok=True)
 
 transform_gallery = transforms.Compose([transforms.ToPILImage()])
-
 test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())
 
-# Pick a few samples (car, dog, plane, cat)
-sample_indices = [1, 3, 10, 25]  # arbitrary indices
+# Collect one sample for each class
 example_images = []
-for idx in sample_indices:
+seen_classes = set()
+for idx in range(len(test_set)):
     img, label = test_set[idx]
-    pil_img = transform_gallery(img)
-    file_path = os.path.join(sample_dir, f"example_{class_labels[label]}.png")
-    pil_img.save(file_path)
-    example_images.append(file_path)
+    if label not in seen_classes:
+        pil_img = transform_gallery(img)
+        file_path = os.path.join(sample_dir, f"example_{class_labels[label]}.png")
+        pil_img.save(file_path)
+        example_images.append(file_path)
+        seen_classes.add(label)
+    if len(seen_classes) == 10:
+        break
 
 # === Gradio Interface Setup ===
 interface = gr.Interface(