Update app.py

app.py

@@ -4,8 +4,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 import numpy as np
+import matplotlib.pyplot as plt
 
-# Define the network (same as your model)
+# Define the network
 class Net(nn.Module):
     def __init__(self):
         super(Net, self).__init__()
@@ -24,88 +25,72 @@ class Net(nn.Module):
         x = self.fc2(x)
         return F.log_softmax(x, dim=1)
 
-# Load the model
+# Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 network = Net().to(device)
 network.load_state_dict(torch.load("model.pth", map_location=device))
 network.eval()
 
-# Prediction function
 def predict_digit(image):
-    """
-    Takes an image from the drawing pad and predicts the digit
-    """
+    """Predict digit from image"""
     if image is None:
-        return "No input", "0%"
+        return None
     
-    # Convert image to grayscale and resize to 28x28
-    if isinstance(image, dict):  # Gradio 6.0 returns dict with 'composite' key
+    # Handle Gradio 6.0 dict format
+    if isinstance(image, dict):
         image = image['composite']
     
     img = Image.fromarray(image).convert('L')
     img = img.resize((28, 28))
     
-    # Convert to tensor
     img_array = np.array(img, dtype=np.float32) / 255.0
     img_tensor = torch.from_numpy(img_array).unsqueeze(0).unsqueeze(0).to(device)
-    
-    # Apply same normalization as training
     img_tensor = (img_tensor - 0.1307) / 0.3081
     
-    # Predict
     with torch.no_grad():
         output = network(img_tensor)
-        probabilities = torch.exp(output[0])
-        pred = output.data.max(1, keepdim=True)[1][0].item()
-        confidence = probabilities[pred].item() * 100
+        probs = torch.exp(output[0])
+        pred = output.data.max(1)[1][0].item()
     
-    # Create confidence text and plot data
-    confidence_text = f"Predicted: **{int(pred)}** ({confidence:.1f}%)"
+    # Create visualization
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
     
-    # Format for bar chart (as list of tuples)
-    chart_data = [(str(i), float(probabilities[i].item() * 100)) for i in range(10)]
+    # Plot 1: Drawn digit
+    ax1.imshow(img, cmap='gray')
+    ax1.set_title(f'Predicted: {int(pred)}', fontsize=16, fontweight='bold')
+    ax1.axis('off')
     
-    return confidence_text, chart_data
-
-# Create Gradio interface
-title = "🔢 MNIST Digit Recognition"
-description = "Draw a digit (0-9) on the canvas and the model will predict what it is!"
+    # Plot 2: Confidence bars
+    digits = list(range(10))
+    confidences = [float(probs[i].item() * 100) for i in range(10)]
+    colors = ['#4CAF50' if i == pred else '#FF6B6B' for i in range(10)]
+    
+    ax2.bar(digits, confidences, color=colors, alpha=0.7)  # ✅ FIXED: Added closing parenthesis
+    ax2.set_xlabel('Digit', fontsize=12)
+    ax2.set_ylabel('Confidence (%)', fontsize=12)
+    ax2.set_title('Confidence Scores', fontsize=12, fontweight='bold')
+    ax2.set_ylim([0, 100])
+    ax2.set_xticks(digits)
+    ax2.grid(axis='y', alpha=0.3)
+    
+    plt.tight_layout()
+    return fig
 
-with gr.Blocks(title=title) as demo:
-    gr.Markdown(f"# {title}")
-    gr.Markdown(description)
+# Create interface
+with gr.Blocks(title="MNIST Digit Recognition") as demo:
+    gr.Markdown("# 🔢 MNIST Digit Recognition")
+    gr.Markdown("Draw a digit (0-9) and the model will predict it!")
     
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### Draw a Digit")
-            canvas = gr.Sketchpad(
-                label="Canvas",
-                type="numpy",
-                interactive=True
-            )
-            with gr.Row():
-                clear_btn = gr.Button("Clear Canvas", size="sm")
-                submit_btn = gr.Button("Predict", variant="primary", size="sm")
-        
-        with gr.Column(scale=1):
-            gr.Markdown("### Prediction Result")
-            output_text = gr.Markdown("Draw a digit to get started!")
-            
-            # Use Plot instead of BarChart for Gradio 6.0
-            confidence_plot = gr.Plot(label="Confidence Scores")
+        canvas = gr.Sketchpad(label="Draw Here", type="numpy")
+        output = gr.Plot(label="Prediction")
     
-    # Function to create bar chart using matplotlib
-    def create_prediction_with_chart(image):
-        pred_text, chart_data = predict_digit(image)
-        
-        # Create matplotlib figure
-        import matplotlib.pyplot as plt
-        fig, ax = plt.subplots(figsize=(8, 4))
-        
-        digits = [str(i) for i in range(10)]
-        confidences = [data[1] for data in chart_data]
-        
-        colors = ['#FF6B6B' if i != int(chart_data[int(chart_data[0][0])][0]) else '#4CAF50' 
-                  for i in range(10)]
-        
-        ax.bar(digits, confidences, color=colors, alpha
+    with gr.Row():
+        predict_btn = gr.Button("Predict", variant="primary")
+        clear_btn = gr.Button("Clear")
+    
+    predict_btn.click(fn=predict_digit, inputs=canvas, outputs=output)
+    clear_btn.click(lambda: None, outputs=canvas)
+
+if __name__ == "__main__":
+    demo.launch()