Update utils.py

utils.py

@@ -82,93 +82,61 @@
 #     top3 = sorted(probs_combined.items(), key=lambda x: x[1], reverse=True)[:3]
 #     return {idx_to_class[k]: float(f"{v:.4f}") for k, v in top3}
 
-
-
 import torch
-import torch.nn.functional as F
-from torchvision import transforms
-from PIL import Image
-import gradio as gr
 import torch.nn as nn
-from torchvision.models import resnet18  # Example; change to your actual architecture
-
-# ✅ Define your 25 class names (index 0 → class 1)
-class_names = [
-    "Capplehinee ", "Lantana", "Negative", "Parkinsonia", "Parthenium", "Prickly acacia",
-    "Rubber vine", "Siam weed", "Snake weed",        # 1-9 (Model 1)
-    "Broadleaf",           # class10 (Model 3)
-    "Grass",               # class11
-    "Soil",                # class12
-    "Soybean",             # class13
-    "Black grass",         # class14 (Model 2)
-    "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"   # 14-25 (Model 2)
-]
-
-# ✅ Define transforms (adjust to match your model training)
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.5, 0.5, 0.5],
-                         std=[0.5, 0.5, 0.5])
-])
-
-# ✅ Load your models (use correct architecture and weights)
-def get_model(num_classes):
-    model = resnet18(pretrained=False)
-    model.fc = nn.Linear(model.fc.in_features, num_classes)
-    return model
-
-model1 = get_model(9)     # for class 1-9
-model3 = get_model(4)     # for class 10-13
-model2 = get_model(12)    # for class 14-25
-
-model1.load_state_dict(torch.load("MMIM_best1.pth", map_location='cpu'))
-model2.load_state_dict(torch.load("MMIM_best2.pth", map_location='cpu'))
-model3.load_state_dict(torch.load("MMIM_best3.pth", map_location='cpu'))
-
-model1.eval()
-model2.eval()
-model3.eval()
-
-# ✅ Inference function
-def predict(image):
-    image_tensor = transform(image).unsqueeze(0)
+from torchvision.models.swin_transformer import swin_t, Swin_T_Weights
+import torch.nn.functional as F
+
+# ✅ Define MMIM architecture (same as used during training)
+class MMIM(nn.Module):
+    def __init__(self, num_classes):
+        super(MMIM, self).__init__()
+        self.backbone = swin_t(weights=Swin_T_Weights.DEFAULT)
+        self.backbone.head = nn.Identity()
+        self.classifier = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        features = self.backbone(x)
+        return self.classifier(features)
+
+# ✅ Load all 3 models
+def load_all_models():
+    model1 = MMIM(num_classes=9)     # class1–9
+    model2 = MMIM(num_classes=12)    # class14–25
+    model3 = MMIM(num_classes=4)     # class10–13
+
+    model1.load_state_dict(torch.load("MMIM_best1.pth", map_location='cpu'))
+    model2.load_state_dict(torch.load("MMIM_best2.pth", map_location='cpu'))
+    model3.load_state_dict(torch.load("MMIM_best3.pth", map_location='cpu'))
+
+    model1.eval()
+    model2.eval()
+    model3.eval()
+
+    return model1, model2, model3
+
+# ✅ Inference combining raw logits before softmax
+def predict_image(image, model1, model2, model3, transform, class_names):
+    image_tensor = transform(image).unsqueeze(0)  # [1, 3, 224, 224]
 
     with torch.no_grad():
-        out1 = F.softmax(model1(image_tensor), dim=1)      # [1, 9]
-        out3 = F.softmax(model3(image_tensor), dim=1)      # [1, 4]
-        out2 = F.softmax(model2(image_tensor), dim=1)      # [1, 12]
-
-    # Combine into a 25-class vector
-    combined = torch.cat([out1, out3, out2], dim=1)  # shape: [1, 25]
-    pred_idx = combined.argmax(dim=1).item()
-    confidence = combined.max().item()
-
-    # Optional rejection
-    if confidence < 0.5:
-        return "Prediction uncertain or unknown class"
-
-    return f"Predicted: {class_names[pred_idx]} (Confidence: {confidence:.2f})"
-
-# ✅ Gradio app
-app = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Weed Classifier - 25 Class Combined (3 Models)",
-    description="Upload an image to classify weeds across 25 species using 3 separate models."
-)
-
-# ✅ Launch
-if __name__ == "__main__":
-    app.launch()
+        logit1 = model1(image_tensor)  # [1, 9]
+        logit3 = model3(image_tensor)  # [1, 4]
+        logit2 = model2(image_tensor)  # [1, 12]
+
+    # ✅ Combine logits (not softmax) → then apply softmax
+    combined_logits = torch.cat([logit1, logit3, logit2], dim=1)  # [1, 25]
+    combined_probs = F.softmax(combined_logits, dim=1)            # unified softmax
+
+    pred_idx = combined_probs.argmax(dim=1).item()
+    confidence = combined_probs[0, pred_idx].item()
+
+    return class_names[pred_idx], confidence
+
+
+