Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 
 # 🔧 Model definition
 class MMIM(nn.Module):
-    def __init__(self, num_classes=9):
+    def __init__(self, num_classes=36):
         super(MMIM, self).__init__()
         self.backbone = swin_t(weights='IMAGENET1K_V1')
         self.backbone.head = nn.Identity()
@@ -24,9 +24,9 @@ class MMIM(nn.Module):
 
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MMIM(num_classes=9)
+model = MMIM(num_classes=36)
 
-# 🧠 Load pretrained weights except mismatched classifier
+# 🧠 Load only matching weights from checkpoint (skip classifier mismatch)
 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
@@ -36,12 +36,16 @@ model.load_state_dict(filtered_checkpoint, strict=False)
 model.to(device)
 model.eval()
 
-# ✅ Updated class names (match folder structure)
+# ✅ Correct class names list (ordered by folder names)
 class_names = [
     "Chinee apple", "Lantana", "Negative", "Parkinsonia", "Parthenium",
-    "Prickly acacia", "Rubber vine", "Siam weed", "Snake weed", 'Black grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common Wheat',
-    'Fat Hen', 'Loose Silky-bent', 'Maize', 'Scentless Mayweed',
-    'Shepherds purse', 'Small-flowered Cranesbill', 'Sugar beet','Carpetweeds','Crabgrass','Eclipta','Goosegrass','Morningglory','Nutsedge','PalmerAmaranth','Pricky Sida','Purslane','Ragweed','Sicklepod','SpottedSpurge','SpurredAnoda','Swinecress','Waterhemp'
+    "Prickly acacia", "Rubber vine", "Siam weed", "Snake weed", 
+    "Black grass", "Charlock", "Cleavers", "Common Chickweed", "Common Wheat",
+    "Fat Hen", "Loose Silky-bent", "Maize", "Scentless Mayweed",
+    "Shepherds purse", "Small-flowered Cranesbill", "Sugar beet",
+    "Carpetweeds", "Crabgrass", "Eclipta", "Goosegrass", "Morningglory",
+    "Nutsedge", "PalmerAmaranth", "Pricky Sida", "Purslane", "Ragweed",
+    "Sicklepod", "SpottedSpurge", "SpurredAnoda", "Swinecress", "Waterhemp"
 ]
 
 # 🔁 Image transform
@@ -50,7 +54,7 @@ transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
-# 🔍 Prediction function with negative detection
+# 🔍 Prediction function
 def predict(img):
     img = img.convert('RGB')
     img_tensor = transform(img).unsqueeze(0).to(device)