Update app.py

app.py

@@ -6,6 +6,7 @@ import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
 import os
+import numpy as np
 
 # === Simple CNN Model Definition ===
 class SimpleCNN(nn.Module):
@@ -45,7 +46,11 @@ preprocess = transforms.Compose([
     transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
 ])
 
-# === Inference Function ===
+# === CIFAR-10 Test Loader for Benchmark Mode ===
+test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())
+test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False)
+
+# === Inference Function (single image) ===
 def inference(input_image: Image.Image):
     if model.training:
         model.eval()
@@ -56,58 +61,53 @@ def inference(input_image: Image.Image):
     confidences = {class_labels[i]: float(probabilities[0,i]) for i in range(len(class_labels))}
     return confidences
 
-# === Results Viewer Function ===
-def show_results(input_image: Image.Image):
-    preds = inference(input_image)
+# === Benchmark Mode: Evaluate on full test set ===
+def benchmark():
+    model.eval()
+    correct = 0
+    total = 0
+    class_correct = np.zeros(10)
+    class_total = np.zeros(10)
+
+    with torch.no_grad():
+        for inputs, labels in test_loader:
+            outputs = model(inputs)
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+            c = (predicted == labels).squeeze()
+            for i in range(len(labels)):
+                label = labels[i].item()
+                class_correct[label] += c[i].item()
+                class_total[label] += 1
+
+    overall_acc = 100.0 * correct / total
+    classwise_acc = {class_labels[i]: round(100.0 * class_correct[i] / class_total[i], 2) for i in range(10)}
 
     # Load plots if they exist
     perf_plot = "training_performance.png" if os.path.exists("training_performance.png") else None
     acc_plot = "final_test_accuracy.png" if os.path.exists("final_test_accuracy.png") else None
 
-    # Load final test accuracy number
-    test_acc_text = "Final test accuracy not available."
-    if os.path.exists("final_test_accuracy.txt"):
-        with open("final_test_accuracy.txt", "r") as f:
-            test_acc_value = f.read().strip()
-        test_acc_text = f"Final Test Accuracy: {test_acc_value}%"
+    return overall_acc, classwise_acc, perf_plot, acc_plot
 
-    return preds, perf_plot, acc_plot, test_acc_text
-
-# === Prepare CIFAR-10 Sample Gallery (one per class) ===
-sample_dir = "examples"
-os.makedirs(sample_dir, exist_ok=True)
+# === Gradio Interface Setup ===
+with gr.Blocks() as demo:
+    gr.Markdown("## CIFAR-10 Image Classification with DCLR Optimizer")
+    gr.Markdown("Upload an image for prediction, or run Benchmark Mode to see full test accuracy.")
 
-transform_gallery = transforms.Compose([transforms.ToPILImage()])
-test_set = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transforms.ToTensor())
+    with gr.Tab("Single Image Inference"):
+        inp = gr.Image(type='pil', label='Upload Image')
+        out = gr.Label(num_top_classes=3, label='Predictions')
+        inp.change(fn=inference, inputs=inp, outputs=out)
 
-# Collect one sample for each class
-example_images = []
-seen_classes = set()
-for idx in range(len(test_set)):
-    img, label = test_set[idx]
-    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
+    with gr.Tab("Benchmark Mode"):
+        btn = gr.Button("Run Benchmark on CIFAR-10 Test Set")
+        overall = gr.Textbox(label="Overall Test Accuracy")
+        classwise = gr.JSON(label="Per-Class Accuracy (%)")
+        perf_plot = gr.Image(type='filepath', label='Training Performance')
+        acc_plot = gr.Image(type='filepath', label='Final Test Accuracy Plot')
 
-# === Gradio Interface Setup ===
-interface = gr.Interface(
-    fn=show_results,
-    inputs=gr.Image(type='pil', label='Upload Image'),
-    outputs=[
-        gr.Label(num_top_classes=3, label='Predictions'),
-        gr.Image(type='filepath', label='Training Performance'),
-        gr.Image(type='filepath', label='Final Test Accuracy Plot'),
-        gr.Textbox(label='Final Test Accuracy')
-    ],
-    title='CIFAR-10 Image Classification with DCLR Optimizer',
-    description='Upload an image or try sample CIFAR-10 images. See predictions plus benchmark plots and accuracy.',
-    examples=example_images
-)
+        btn.click(fn=benchmark, inputs=None, outputs=[overall, classwise, perf_plot, acc_plot])
 
 if __name__ == '__main__':
-    interface.launch()
+    demo.launch()