Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 
 # 🔧 Model definition
 class MMIM(nn.Module):
-    def __init__(self, num_classes=36):
+    def __init__(self, num_classes=40):
         super(MMIM, self).__init__()
         self.backbone = swin_t(weights='IMAGENET1K_V1')
         self.backbone.head = nn.Identity()
@@ -24,56 +24,55 @@ class MMIM(nn.Module):
 
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MMIM(num_classes=36)
+model = MMIM(num_classes=40)
 
-# 🧠 Load only matching weights from checkpoint (skip classifier mismatch)
+# 🧠 Load checkpoint (ignore mismatched keys like classifier weights)
 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()
 
-# ✅ class_names mapped according to confusion matrix order
+# ✅ Corrected class_names mapping based on your confusion matrix and folder order
 class_names = [
     "Chinee apple",              # class1
     "Lantana",                   # class2
     "Negative",                  # class3
-     "Rubber vine",               # class7
-     "Snake weed",                # class9
-     "Black grass",               # class10
-     "Charlock",                  # class11
-      "Cleavers",                  # class12
-      "Common Chickweed"           # class13
-      "Common Wheat",              # class14
-    "Fat Hen",                   # class15
-     "Loose Silky-bent",          # class16
-    "Maize",                     # class17
-    "Scentless Mayweed",         # class18
-    "Shepherds purse",           # class19
-    "Small-flowered Cranesbill",# class20
-    "Sugar beet",                # class21
-    "Carpetweeds",               # class22
-    "Crabgrass",                 # class23
-    "Eclipta",                   # class24
-    "Goosegrass",                # class25
-    "Morningglory",              # class26
-    "Nutsedge",                  # class27
-    "PalmerAmaranth",            # class28
-    "Pricky Sida",               # class29
-    "Purslane",                  # class30
-    "Ragweed",                   # class31
-    "Sicklepod",                 # class32
-    "SpottedSpurge",             # class33
-    "SpurredAnoda",              # class34
-    "Swinecress",                # class35
-    "Waterhemp",                 # class36
-    "Parkinsonia",               # class4
-    "Parthenium",                # class5
-    "Prickly acacia",            # class6
-   "Siam weed",                 # class8
+    "Parkinsonia",              # class4
+    "Parthenium",               # class5
+    "Prickly acacia",           # class6
+    "Rubber vine",              # class7
+    "Siam weed",                # class8
+    "Snake weed",               # class9
+    "Black grass",              # class14
+    "Charlock",                 # class15
+    "Cleavers",                 # class16
+    "Common Chickweed",         # class17
+    "Common Wheat",             # class18
+    "Fat Hen",                  # class19
+    "Loose Silky-bent",         # class20
+    "Maize",                    # class21
+    "Scentless Mayweed",        # class22
+    "Shepherds purse",          # class23
+    "Small-flowered Cranesbill",# class24
+    "Sugar beet",               # class25
+    "Carpetweeds",              # class26
+    "Crabgrass",                # class27
+    "Eclipta",                  # class28
+    "Goosegrass",               # class29
+    "Morningglory",             # class30
+    "Nutsedge",                 # class31
+    "PalmerAmaranth",           # class32
+    "Pricky Sida",              # class33
+    "Purslane",                 # class34
+    "Ragweed",                  # class35
+    "Sicklepod",                # class36
+    "SpottedSpurge",            # class37
+    "Spurred Anoda",            # class38
+    "Swinecress",               # class39
+    "Waterhemp"                 # class40
 ]
 
 # 🔁 Image transform