Create app.py

app.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from torchvision import transforms
+import gradio as gr
+from PIL import Image
+
+# 1. Define the EXACT same architecture from your notebook
+class VehicleClassifier(nn.Module):
+    def __init__(self):
+        super(VehicleClassifier, self).__init__()
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2)
+        )
+        self.fc_layers = nn.Sequential(
+            nn.Linear(64 * 28 * 28, 512),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(512, 8)
+        )
+
+    def forward(self, x):
+        x = self.conv_layers(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc_layers(x)
+        return x
+
+# 2. Setup Device and Load Model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = VehicleClassifier().to(device)
+
+# Replace with your actual repo and filename
+path = "abhiprd20/vehicle-classification-utd" 
+model.load_state_dict(torch.hub.load_state_dict_from_url(
+    f"https://huggingface.co/{path}/resolve/main/model.pth", 
+    map_location=device
+))
+model.eval()
+
+# 3. Prediction Logic
+classes = ['Bicycle', 'Bus', 'Car', 'Motorcycle', 'NonVehicles', 'Taxi', 'Truck', 'Van']
+data_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def predict(img):
+    img = data_transforms(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        outputs = model(img)
+        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+        confidences = {classes[i]: float(probs[i]) for i in range(8)}
+    return confidences
+
+# 4. Interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    title="Autonomous Vehicle Perception Demo",
+    description="Custom CNN Model for UT Dallas AI Safety Lab. Target: 78.54% Accuracy."
+).launch()