Update app.py

app.py

@@ -1,76 +1,65 @@
-import gradio as gr
-import torch
-import torch.nn as nn
-from torchvision import transforms, models
-from PIL import Image
-
-# ===============================
-# 📌 Classes (same order as ImageFolder during training)
-# Replace with trainers.classes if available
-# ===============================
-class_names = [
-    "Acura", "Alfa Romeo", "Aston Martin", "Audi", "Bentley",
-    "BMW", "Bugatti", "Buick", "Cadillac", "Chevrolet",
-    "Chrysler", "Citroen", "Dodge", "Ferrari", "Fiat",
-    "Ford", "Genesis", "GMC", "Honda", "Hyundai",
-    "Infiniti", "Jaguar", "Jeep", "Kia", "Lamborghini",
-    "Land Rover", "Lexus", "Lincoln", "Maserati", "Mazda",
-    "McLaren", "Mercedes", "Mini", "Mitsubishi", "Nissan",
-    "Pagani", "Peugeot", "Porsche", "Ram", "Renault",
-    "Rolls Royce", "Saab", "Subaru", "Suzuki", "Tesla",
-    "Toyota", "Volkswagen", "Volvo", "Others1", "Others2"
-]
-
-# ===============================
-# 📌 Inference Transform (NO augmentation)
-# ===============================
-transform = transforms.Compose([
-    transforms.Lambda(lambda x: x.convert("RGB")),
-    transforms.Resize((224,224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.5]*3, [0.5]*3)
-])
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ===============================
-# 📌 Load Model
-# ===============================
-def load_model(weights_path="best_model.pth"):
-    base_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
-    in_features = base_model.fc.in_features
-    base_model.fc = nn.Sequential(
-        nn.Linear(in_features, 512),
-        nn.ReLU(),
-        nn.Dropout(0.5),
-        nn.Linear(512, len(class_names))
-    )
-    base_model.load_state_dict(torch.load(weights_path, map_location=device))
-    return base_model.to(device).eval()
-
-model = load_model()
-
-# ===============================
-# 📌 Prediction Function
-# ===============================
-def predict(image):
-    img = transform(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        outputs = model(img)
-        probs = torch.softmax(outputs, dim=1)[0]
-        top5_probs, top5_idx = torch.topk(probs, 5)
-    return {class_names[idx]: float(prob) for idx, prob in zip(top5_idx, top5_probs)}
-
-# ===============================
-# 📌 Gradio Interface
-# ===============================
-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 and the model predicts the brand (Top-5)."
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import torch
+import torch.nn as nn
+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_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
+}
+idx_to_class = {v: k for k, v in class_to_idx.items()}
+
+# -----------------------------
+# 3. Transform (inference version: no randomness)
+# -----------------------------
+transform = transforms.Compose([
+    transforms.Lambda(lambda x: x.convert("RGB")),
+    transforms.Resize((224,224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.5]*3, [0.5]*3)
+])
+
+# -----------------------------
+# 4. Prediction function
+# -----------------------------
+def predict(img):
+    img = transform(img).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = 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
+
+# -----------------------------
+# 5. 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"
+)
+
+if __name__ == "__main__":
+    demo.launch()