Update app.py

app.py

@@ -1,94 +1,49 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-from torchvision import transforms
+from torchvision import models, transforms
 from PIL import Image
 
-# ---------------------------------------------------------
-# 1. MODEL ARCHITECTURE
-# ---------------------------------------------------------
-class SimpleCNN(nn.Module):
-    def __init__(self, num_classes=10):
-        super(SimpleCNN, self).__init__()
-        
-        self.conv_block1 = nn.Sequential(
-            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-        )
+# 1. SETUP MODEL
+# We use ResNet18 structure to match your training
+model = models.resnet18(weights=None) 
+model.fc = nn.Linear(model.fc.in_features, 10) # Adjust head to 10 classes
 
-        self.conv_block2 = nn.Sequential(
-            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-        )
-
-        self.classifier = nn.Sequential(
-            nn.Flatten(),
-            nn.Linear(in_features=32 * 32 * 32, out_features=128),
-            nn.ReLU(),
-            nn.Linear(in_features=128, out_features=num_classes),
-        )
-
-    def forward(self, x):
-        x = self.conv_block1(x)
-        x = self.conv_block2(x)
-        x = self.classifier(x)
-        return x
-
-# ---------------------------------------------------------
-# 2. SETUP
-# ---------------------------------------------------------
-# Initialize model
-model = SimpleCNN()
-
-# Load weights (Ensure 'fulldigits.pt' is uploaded to Hugging Face Files!)
+# Load your 98.79% accuracy weights
 try:
-    model.load_state_dict(torch.load("fulldigits.pt", map_location="cpu"))
+    state_dict = torch.load("fulldigits.pt", map_location="cpu")
+    model.load_state_dict(state_dict)
     model.eval()
-except FileNotFoundError:
-    print("Error: 'fulldigits.pt' not found. Please upload your model file.")
+except Exception as e:
+    print(f"Error loading model: {e}")
 
-# Define transforms
-# CRITICAL FIX: Added lambda to force RGB. 
-# This prevents crashes if someone uploads a Grayscale or RGBA image.
+# 2. PREPROCESSING
+# Must use the ImageNet stats you trained with!
 transform = transforms.Compose([
-    transforms.Lambda(lambda x: x.convert("RGB")), 
-    transforms.Resize((128, 128)),
+    transforms.Lambda(lambda x: x.convert("RGB")), # Force RGB
+    transforms.Resize((128, 128)),                 # Match training size
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# ---------------------------------------------------------
-# 3. PREDICTION FUNCTION
-# ---------------------------------------------------------
+# 3. PREDICT FUNCTION
 def predict(image):
-    if image is None:
-        return None
-
-    # Transform image
+    if image is None: return None
     img_tensor = transform(image).unsqueeze(0)
-
-    # Make prediction
+    
     with torch.no_grad():
         output = model(img_tensor)
-        # Get probabilities
         probabilities = torch.nn.functional.softmax(output[0], dim=0)
     
-    # Return a dictionary for Gradio's Label component
-    # This creates the nice bar chart effect
     return {str(i): float(probabilities[i]) for i in range(10)}
 
-# ---------------------------------------------------------
-# 4. GRADIO INTERFACE
-# ---------------------------------------------------------
+# 4. INTERFACE
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Label(num_top_classes=3, label="Predictions"), # Changed to Label for better UI
-    title="Digit Classification Project",
-    description="Upload an image to check if it contains a digit (0-9).",
-    # removed share=True for production deployment
+    inputs=gr.Image(type="pil", label="Draw or Upload Digit"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Handwritten Digit Recognizer",
+    description="A ResNet18 model fine-tuned to 98.79% accuracy."
 )
 
 if __name__ == "__main__":