Update app.py

app.py

@@ -25,14 +25,21 @@ class MMIM(nn.Module):
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = MMIM(num_classes=9)
-model.load_state_dict(torch.load("MMIM_best.pth", map_location=device))
+
+# 🧠 Load pretrained weights except mismatched classifier
+checkpoint = torch.load("MMIM_best.pth", map_location=device)
+filtered_checkpoint = {
+    k: v for k, v in checkpoint.items() if k in model.state_dict() and model.state_dict()[k].shape == v.shape
+}
+model.load_state_dict(filtered_checkpoint, strict=False)
+
 model.to(device)
 model.eval()
 
 # ✅ Updated class names (match folder structure)
 class_names = [
     "Chinee apple", "Lantana", "Negative", "Parkinsonia", "Parthenium",
-    "Prickly acacia", "Rubber vine", "Siam weed", "Snake weed"
+    "Prickly acacia", "Rubber vine", "Siam weed", "Snake weed",
 ]
 
 # 🔁 Image transform