Update app.py

app.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 import gradio as gr
-import os # Import os to check for model file
+import os
 
 # === Simple CNN Model Definition ===
 class SimpleCNN(nn.Module):
@@ -27,67 +27,66 @@ class SimpleCNN(nn.Module):
 model = SimpleCNN()
 model_path = 'simple_cnn_dclr_tuned.pth'
 
-# Check if the model file exists before loading
 if os.path.exists(model_path):
     model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-    model.eval() # Set model to evaluation mode
+    model.eval()
     print(f"Model loaded successfully from {model_path}")
 else:
-    print(f"Warning: Model file '{model_path}' not found. Please ensure 'train_dclr_model.py' has been run.")
-    # Optionally, you might want to exit or raise an error if the model is crucial
-
+    print(f"Warning: Model file '{model_path}' not found. Please run train_dclr_model.py first.")
 
 # === CIFAR-10 Class Labels ===
-class_labels = ['plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+class_labels = ['plane','car','bird','cat','deer','dog','frog','horse','ship','truck']
 
 # === Image Preprocessing ===
 preprocess = transforms.Compose([
-    transforms.Resize(32), # CIFAR-10 images are 32x32
+    transforms.Resize(32),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet stats are common
+    transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])
 ])
 
 # === Inference Function ===
 def inference(input_image: Image.Image):
-    if model.training: # Ensure model is in eval mode
+    if model.training:
         model.eval()
-
-    # Preprocess the image
-    processed_image = preprocess(input_image)
-    # Add a batch dimension
-    processed_image = processed_image.unsqueeze(0)
-
-    # Perform inference
+    processed_image = preprocess(input_image).unsqueeze(0)
     with torch.no_grad():
         outputs = model(processed_image)
         probabilities = F.softmax(outputs, dim=1)
-
-    # Convert probabilities to a dictionary of class labels and scores
-    confidences = {class_labels[i]: float(probabilities[0, i]) for i in range(len(class_labels))}
+    confidences = {class_labels[i]: float(probabilities[0,i]) for i in range(len(class_labels))}
     return confidences
 
-# === Gradio Interface Setup ===
-# Example images (replace with actual paths if available, or keep as dummy for now)
-# For a Hugging Face Space, you might place example images in an 'examples/' directory.
-example_images = [
-    # os.path.join(os.path.dirname(__file__), "examples/example_car.png"),
-    # os.path.join(os.path.dirname(__file__), "examples/example_dog.png"),
-    # os.path.join(os.path.dirname(__file__), "examples/example_plane.png")
-]
+# === Results Viewer Function ===
+def show_results(input_image: Image.Image):
+    preds = inference(input_image)
 
-# A placeholder for example images since we don't have them generated yet.
-# Users can upload their own or I will add some placeholder images if needed in the next step.
-# For now, an empty list of examples is fine.
+    # 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 if available
+    test_acc_text = "Final Test Accuracy plot not found."
+    if acc_plot and os.path.exists("final_test_accuracy.png"):
+        # You can optionally parse accuracy from a saved file; here we just show a placeholder
+        test_acc_text = "See bar chart above for final test accuracy."
+
+    return preds, perf_plot, acc_plot, test_acc_text
+
+# === Gradio Interface Setup ===
+example_images = []
 
 interface = gr.Interface(
-    fn=inference,
-    inputs=gr.Image(type='pil', label='Input Image'),
-    outputs=gr.Label(num_top_classes=3, label='Predictions'),
+    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 and see the model\'s predictions using a SimpleCNN trained with the DCLR optimizer.',
+    description='Upload an image to see predictions. Training/test plots show benchmark results on CIFAR-10.',
     examples=example_images
 )
 
-# === Launch Gradio App ===
 if __name__ == '__main__':
     interface.launch()