Update app.py

app.py

@@ -3,13 +3,11 @@ 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):
-    def __init__(self, num_classes=40):
+    def __init__(self, num_classes=36):
         super(MMIM, self).__init__()
         self.backbone = swin_t(weights='IMAGENET1K_V1')
         self.backbone.head = nn.Identity()
@@ -26,7 +24,7 @@ class MMIM(nn.Module):
 
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MMIM(num_classes=40)
+model = MMIM(num_classes=36)
 
 checkpoint = torch.load("MMIM_best.pth", map_location=device)
 filtered_checkpoint = {
@@ -37,57 +35,46 @@ model.load_state_dict(filtered_checkpoint, strict=False)
 model.to(device)
 model.eval()
 
-# ✅ 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))]
+# ✅ Final class_names (corrected index order, without class10–class13)
+class_names = [
+    "Chinee apple",              # 0 = class1
+    "Lantana",                   # 1 = class2
+    "Negative",                  # 2 = class3
+    "Parkinsonia",              # 3 = class4
+    "Parthenium",               # 4 = class5
+    "Prickly acacia",           # 5 = class6
+    "Rubber vine",              # 6 = class7
+    "Siam weed",                # 7 = class8
+    "Snake weed",               # 8 = class9
+    # Skipping class10–class13
+    "Common Wheat",             # 9 = class14
+    "Fat Hen",                  # 10 = class15
+    "Loose Silky-bent",         # 11 = class16
+    "Maize",                    # 12 = class17
+    "Scentless Mayweed",        # 13 = class18
+    "Shepherds purse",          # 14 = class19
+    "Small-flowered Cranesbill",# 15 = class20
+    "Sugar beet",               # 16 = class21
+    "Carpetweeds",              # 17 = class22
+    "Crabgrass",                # 18 = class23
+    "Eclipta",                  # 19 = class24
+    "Goosegrass",               # 20 = class25
+    "Morningglory",             # 21 = class26
+    "Nutsedge",                 # 22 = class27
+    "PalmerAmaranth",           # 23 = class28
+    "Pricky Sida",              # 24 = class29
+    "Purslane",                 # 25 = class30
+    "Ragweed",                  # 26 = class31
+    "Sicklepod",                # 27 = class32
+    "SpottedSpurge",            # 28 = class33
+    "Spurred Anoda",            # 29 = class34
+    "Swinecress",               # 30 = class35
+    "Waterhemp",                # 31 = class36
+    "Extra1",                   # 32 = class37 (if any)
+    "Extra2",                   # 33 = class38
+    "Extra3",                   # 34 = class39
+    "Extra4"                    # 35 = class40
+]
 
 # 🔁 Image transform
 transform = transforms.Compose([