Update app.py

app.py

@@ -1,5 +1,6 @@
 import torch
 import torchvision.transforms as transforms
+import torchvision
 import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
@@ -72,9 +73,26 @@ def show_results(input_image: Image.Image):
 
     return preds, perf_plot, acc_plot, test_acc_text
 
-# === Gradio Interface Setup ===
+# === Prepare CIFAR-10 Sample Gallery ===
+# Download CIFAR-10 test set and save a few sample images
+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
 example_images = []
+for idx in sample_indices:
+    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)
 
+# === Gradio Interface Setup ===
 interface = gr.Interface(
     fn=show_results,
     inputs=gr.Image(type='pil', label='Upload Image'),
@@ -85,7 +103,7 @@ interface = gr.Interface(
         gr.Textbox(label='Final Test Accuracy')
     ],
     title='CIFAR-10 Image Classification with DCLR Optimizer',
-    description='Upload an image to see predictions. Training/test plots and accuracy show benchmark results on CIFAR-10.',
+    description='Upload an image or try sample CIFAR-10 images. See predictions plus benchmark plots and accuracy.',
     examples=example_images
 )