Update utils.py

utils.py

@@ -1,83 +1,174 @@
+# import torch
+# import torch.nn as nn
+# import torch.nn.functional as F
+# from torchvision.models import swin_t
+# from torchvision import transforms
+# from PIL import Image
+# import os
+
+# # --- MMIM model class ---
+# class MMIM(nn.Module):
+#     def __init__(self, num_classes):
+#         super(MMIM, self).__init__()
+#         self.backbone = swin_t(weights='IMAGENET1K_V1')
+#         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):
+#         x = self.backbone(x)
+#         return self.classifier(x)
+
+# # --- Load models with offsets ---
+# def load_all_models():
+#     model_defs = [
+#         ("MMIM_best1.pth", 9),
+#         ("MMIM_best3.pth", 4),
+#         ("MMIM_best2.pth", 12)
+#     ]
+#     device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+#     models = []
+#     offsets = []
+#     total_classes = 0
+#     for path, num_classes in model_defs:
+#         model = MMIM(num_classes)
+#         state_dict = torch.load(path, map_location=device)
+#         model.load_state_dict(state_dict)
+#         model.to(device)
+#         model.eval()
+#         models.append(model)
+#         offsets.append(total_classes)
+#         total_classes += num_classes
+
+#     # Generate dummy class names like class0, class1, ...
+#     idx_to_class = {i: f"class{i}" for i in range(total_classes)}
+#     return models, offsets, idx_to_class
+
+# # --- Inference on one image ---
+# def predict_image(image, models, offsets, idx_to_class):
+#     device = 'cuda' if torch.cuda.is_available() else 'cpu'
+#     transform = transforms.Compose([
+#         transforms.Resize((224, 224)),
+#         transforms.ToTensor(),
+#         transforms.Normalize([0.5]*3, [0.5]*3)
+#     ])
+#     image_tensor = transform(image).unsqueeze(0).to(device)
+
+#     temperatures = [1.2, 1.0, 0.8]  # Adjust for balancing confidence
+
+#     max_score = float('-inf')
+#     final_pred = -1
+#     probs_combined = {}
+
+#     for model, offset, temp in zip(models, offsets, temperatures):
+#         with torch.no_grad():
+#             logits = model(image_tensor) / temp
+#             probs = F.softmax(logits, dim=1).squeeze(0)
+#             top_score, top_class = torch.max(probs, dim=0)
+#             if top_score.item() > max_score:
+#                 max_score = top_score.item()
+#                 final_pred = top_class.item() + offset
+
+#             # Also collect probabilities for all classes
+#             for i, p in enumerate(probs):
+#                 probs_combined[offset + i] = p.item()
+
+#     # Sort top 3
+#     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 as nn
 import torch.nn.functional as F
-from torchvision.models import swin_t
 from torchvision import transforms
 from PIL import Image
-import os
-
-# --- MMIM model class ---
-class MMIM(nn.Module):
-    def __init__(self, num_classes):
-        super(MMIM, self).__init__()
-        self.backbone = swin_t(weights='IMAGENET1K_V1')
-        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):
-        x = self.backbone(x)
-        return self.classifier(x)
-
-# --- Load models with offsets ---
-def load_all_models():
-    model_defs = [
-        ("MMIM_best1.pth", 9),
-        ("MMIM_best3.pth", 4),
-        ("MMIM_best2.pth", 12)
-    ]
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-    models = []
-    offsets = []
-    total_classes = 0
-    for path, num_classes in model_defs:
-        model = MMIM(num_classes)
-        state_dict = torch.load(path, map_location=device)
-        model.load_state_dict(state_dict)
-        model.to(device)
-        model.eval()
-        models.append(model)
-        offsets.append(total_classes)
-        total_classes += num_classes
-
-    # Generate dummy class names like class0, class1, ...
-    idx_to_class = {i: f"class{i}" for i in range(total_classes)}
-    return models, offsets, idx_to_class
-
-# --- Inference on one image ---
-def predict_image(image, models, offsets, idx_to_class):
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.5]*3, [0.5]*3)
-    ])
-    image_tensor = transform(image).unsqueeze(0).to(device)
-
-    temperatures = [1.2, 1.0, 0.8]  # Adjust for balancing confidence
-
-    max_score = float('-inf')
-    final_pred = -1
-    probs_combined = {}
-
-    for model, offset, temp in zip(models, offsets, temperatures):
-        with torch.no_grad():
-            logits = model(image_tensor) / temp
-            probs = F.softmax(logits, dim=1).squeeze(0)
-            top_score, top_class = torch.max(probs, dim=0)
-            if top_score.item() > max_score:
-                max_score = top_score.item()
-                final_pred = top_class.item() + offset
-
-            # Also collect probabilities for all classes
-            for i, p in enumerate(probs):
-                probs_combined[offset + i] = p.item()
-
-    # Sort top 3
-    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 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)
+
+    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()