Update app.py

app.py

@@ -1,64 +1,89 @@
 import torch
 import torch.nn as nn
+import torchvision.models as models
 import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
 
-# -----------------------------
-# 1. Load your model
-# -----------------------------
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = torch.load("best_model.pth", map_location=device)  # update filename if needed
-model.eval()
-
-# -----------------------------
-# 2. Class mapping (from training)
-# -----------------------------
+# --------------------
+# Class Mapping
+# --------------------
 class_to_idx = {
-    'Acura': 0, 'Alfa Romeo': 1, 'Aston Martin': 2, 'Audi': 3, 'BMW': 4, 'Bentley': 5, 'Bugatti': 6,
-    'Buick': 7, 'Cadillac': 8, 'Chevrolet': 9, 'Chrysler': 10, 'Citroen': 11, 'Daewoo': 12,
-    'Dodge': 13, 'Ferrari': 14, 'Fiat': 15, 'Ford': 16, 'GMC': 17, 'Genesis': 18, 'Honda': 19,
-    'Hudson': 20, 'Hyundai': 21, 'Infiniti': 22, 'Jaguar': 23, 'Jeep': 24, 'Kia': 25, 'Land Rover': 26,
-    'Lexus': 27, 'Lincoln': 28, 'MG': 29, 'Maserati': 30, 'Mazda': 31, 'Mercedes-Benz': 32, 'Mini': 33,
-    'Mitsubishi': 34, 'Nissan': 35, 'Oldsmobile': 36, 'Peugeot': 37, 'Pontiac': 38, 'Porsche': 39,
-    'Ram Trucks': 40, 'Renault': 41, 'Saab': 42, 'Studebaker': 43, 'Subaru': 44, 'Suzuki': 45,
-    'Tesla': 46, 'Toyota': 47, 'Volkswagen': 48, 'Volvo': 49
+    'Acura': 0, 'Alfa Romeo': 1, 'Aston Martin': 2, 'Audi': 3, 'BMW': 4,
+    'Bentley': 5, 'Bugatti': 6, 'Buick': 7, 'Cadillac': 8, 'Chevrolet': 9,
+    'Chrysler': 10, 'Citroen': 11, 'Daewoo': 12, 'Dodge': 13, 'Ferrari': 14,
+    'Fiat': 15, 'Ford': 16, 'GMC': 17, 'Genesis': 18, 'Honda': 19,
+    'Hudson': 20, 'Hyundai': 21, 'Infiniti': 22, 'Jaguar': 23, 'Jeep': 24,
+    'Kia': 25, 'Land Rover': 26, 'Lexus': 27, 'Lincoln': 28, 'MG': 29,
+    'Maserati': 30, 'Mazda': 31, 'Mercedes-Benz': 32, 'Mini': 33,
+    'Mitsubishi': 34, 'Nissan': 35, 'Oldsmobile': 36, 'Peugeot': 37,
+    'Pontiac': 38, 'Porsche': 39, 'Ram Trucks': 40, 'Renault': 41,
+    'Saab': 42, 'Studebaker': 43, 'Subaru': 44, 'Suzuki': 45, 'Tesla': 46,
+    'Toyota': 47, 'Volkswagen': 48, 'Volvo': 49
 }
 idx_to_class = {v: k for k, v in class_to_idx.items()}
 
-# -----------------------------
-# 3. Transform (inference version: no randomness)
-# -----------------------------
+# --------------------
+# Image Transform
+# --------------------
 transform = transforms.Compose([
     transforms.Lambda(lambda x: x.convert("RGB")),
     transforms.Resize((224,224)),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomVerticalFlip(p=0.2),
+    transforms.RandomRotation(20),
+    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3),
+    transforms.RandomResizedCrop(224, scale=(0.7, 1.0)),
     transforms.ToTensor(),
     transforms.Normalize([0.5]*3, [0.5]*3)
 ])
 
-# -----------------------------
-# 4. Prediction function
-# -----------------------------
+# --------------------
+# Load Model
+# --------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load pretrained ResNet50 correctly
+base_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+
+# Replace final fully connected layer (your head)
+in_features = base_model.fc.in_features  # 2048 for ResNet50
+base_model.fc = nn.Sequential(
+    nn.Linear(in_features, 512),
+    nn.ReLU(),
+    nn.Dropout(0.5),
+    nn.Linear(512, 50)  # 50 classes
+)
+
+# Load state dict
+state_dict = torch.load("best_model.pth", map_location=device)
+base_model.load_state_dict(state_dict)
+base_model = base_model.to(device)
+base_model.eval()
+
+# --------------------
+# Prediction Function
+# --------------------
 def predict(img):
     img = transform(img).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        outputs = model(img)
+        outputs = base_model(img)
         probs = torch.softmax(outputs, dim=1)[0]
 
-    top5_probs, top5_idx = torch.topk(probs, 5)
-    results = {idx_to_class[idx.item()]: float(top5_probs[i]) for i, idx in enumerate(top5_idx)}
-    return results
+    top5_prob, top5_idx = torch.topk(probs, 5)
+    result = {idx_to_class[idx.item()]: float(top5_prob[i]) for i, idx in enumerate(top5_idx)}
+    return result
 
-# -----------------------------
-# 5. Gradio UI
-# -----------------------------
+# --------------------
+# Gradio UI
+# --------------------
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=5),
     title="Car Brand Classifier",
-    description="Upload a car image to classify its brand"
+    description="Upload a car image to predict its brand."
 )
 
 if __name__ == "__main__":