Update app.py

app.py

@@ -3,7 +3,9 @@ import torch
 import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import swin_t
+from torchvision.datasets import ImageFolder
 from PIL import Image
+import os
 
 # 🔧 Model definition
 class MMIM(nn.Module):
@@ -26,54 +28,66 @@ class MMIM(nn.Module):
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = MMIM(num_classes=40)
 
-# 🧠 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
+    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()
 
-# ✅ Corrected class_names mapping based on your confusion matrix and folder order
-class_names = [
-    "Chinee apple",              # class1
-    "Lantana",                   # class2
-    "Negative",                  # class3
-    "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
-]
+# ✅ Load actual class order from test folder
+test_dir = "test"  # path to the test folder you used with ImageFolder
+folder_class_map = ImageFolder(test_dir).class_to_idx
+idx_to_folder = {v: k for k, v in folder_class_map.items()}
+
+# ✅ Map folder name like 'class15' → weed name
+label_translation = {
+    'class1': "Chinee apple",
+    'class2': "Lantana",
+    'class3': "Negative",
+    'class4': "Parkinsonia",
+    'class5': "Parthenium",
+    'class6': "Prickly acacia",
+    'class7': "Rubber vine",
+    'class8': "Siam weed",
+    'class9': "Snake weed",
+    'class10': "Black grass",
+    'class11': "Charlock",
+    'class12': "Cleavers",
+    'class13': "Common Chickweed",
+    'class14': "Common Wheat",
+    'class15': "Fat Hen",
+    'class16': "Loose Silky-bent",
+    'class17': "Maize",
+    'class18': "Scentless Mayweed",
+    'class19': "Shepherds purse",
+    'class20': "Small-flowered Cranesbill",
+    'class21': "Sugar beet",
+    'class22': "Carpetweeds",
+    'class23': "Crabgrass",
+    'class24': "Eclipta",
+    'class25': "Goosegrass",
+    'class26': "Morningglory",
+    'class27': "Nutsedge",
+    'class28': "PalmerAmaranth",
+    'class29': "Pricky Sida",
+    'class30': "Purslane",
+    'class31': "Ragweed",
+    'class32': "Sicklepod",
+    'class33': "SpottedSpurge",
+    'class34': "Spurred Anoda",
+    'class35': "Swinecress",
+    'class36': "Waterhemp",
+    'class37': "Extra1",
+    'class38': "Extra2",
+    'class39': "Extra3",
+    'class40': "Extra4"
+}
+
+# ✅ Final class_names list (aligned to model output indices)
+class_names = [label_translation[idx_to_folder[i]] for i in range(len(idx_to_folder))]
 
 # 🔁 Image transform
 transform = transforms.Compose([