Create app.py

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import 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),
+        )
+
+        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!)
+try:
+    model.load_state_dict(torch.load("fulldigits.pt", map_location="cpu"))
+    model.eval()
+except FileNotFoundError:
+    print("Error: 'fulldigits.pt' not found. Please upload your model file.")
+
+# Define transforms
+# CRITICAL FIX: Added lambda to force RGB. 
+# This prevents crashes if someone uploads a Grayscale or RGBA image.
+transform = transforms.Compose([
+    transforms.Lambda(lambda x: x.convert("RGB")), 
+    transforms.Resize((128, 128)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+])
+
+# ---------------------------------------------------------
+# 3. PREDICTION FUNCTION
+# ---------------------------------------------------------
+def predict(image):
+    if image is None:
+        return None
+
+    # Transform image
+    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
+# ---------------------------------------------------------
+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
+)
+
+if __name__ == "__main__":
+    demo.launch()