Create app.py

app.py

@@ -0,0 +1,90 @@
+import gradio as gr
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torchvision import transforms
+
+class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+
+device = torch.device("cpu")
+inference_transform = T.Compose([
+    T.ToTensor(),
+    T.Normalize(mean=(0.4914, 0.4822, 0.4465),
+                std=(0.2023, 0.1994, 0.2010)),
+])
+
+class SmallCifarCNN(nn.Module):
+    def __init__(self, num_classes: int = 10):
+        super().__init__()
+        self.features = nn.Sequential(
+            # Block 1
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 32 -> 16
+            # Block 2
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 16 -> 8
+            # Block 3
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2),  # 8 -> 4
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128 * 4 * 4, 256),
+            nn.ReLU(inplace=True),
+            nn.Dropout(p=0.5),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        return x
+
+deployed_model = SmallCifarCNN(num_classes=len(class_names)).to(device)
+
+model_path = 'cifar_cnn_best.pt'
+deployed_model.load_state_dict(
+    torch.load(model_path, map_location=device)
+)
+deployed_model.to(device)
+deployed_model.eval()
+
+def predict_cifar_image(img: Image.Image):
+    """
+    Gradio callback:
+    - Takes a PIL Image
+    - Resizes to 32x32 (CIFAR size)
+    - Normalizes and runs through the CNN
+    - Returns top-3 class probabilities
+    """
+    img = img.convert("RGB")
+    img = img.resize((32, 32), Image.BILINEAR)
+
+    x = inference_transform(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = deployed_model(x)
+        probs = F.softmax(logits, dim=1).cpu().numpy().ravel()
+
+    topk = 3
+    idxs = np.argsort(-probs)[:topk]
+    return {class_names[i]: float(probs[i]) for i in idxs}
+
+demo = gr.Interface(
+    fn=predict_cifar_image,
+    inputs=gr.Image(type="pil", label="Upload an RGB image (will be resized to 32×32)"),
+    outputs=gr.Label(num_top_classes=3, label="Top-3 CIFAR-10 predictions"),
+    title="CIFAR-10 CNN Classifier",
+    description="Small CNN trained on CIFAR-10. Upload an image and see top-3 class probabilities.",
+)
+
+
+if __name__ == "__main__":
+    demo.launch()