Update app.py

app.py

@@ -1,52 +1,21 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from torchvision.models.swin_transformer import swin_t, Swin_T_Weights
 from torchvision import transforms
 from PIL import Image
 import gradio as gr
+from utils import MMIM, load_all_models, predict_image
 
-# ✅ Class labels (indexed 0 to 24 → class1 to class25)
+# ✅ All 25 class labels
 class_names = [
     "Chinee apple", "Lantana", "Negative", "Parkinsonia", "Parthenium", "Prickly acacia",
-    "Rubber vine", "Siam weed", "Snake weed",      # Model 1 (1–9)
-    "Broadleaf", "Grass", "Soil", "Soybean",       # Model 3 (10–13)
+    "Rubber vine", "Siam weed", "Snake weed",      # 1–9 (model1)
+    "Broadleaf", "Grass", "Soil", "Soybean",       # 10–13 (model3)
     "Black grass", "Charlock", "Cleavers", "Common Chickweed", "Common Wheat", "Fat Hen",
     "Loose Silky-bent", "Maize", "Scentless Mayweed", "Shepherds purse",
-    "Small-flowered Cranesbill", "Sugar beet"      # Model 2 (14–25)
+    "Small-flowered Cranesbill", "Sugar beet"      # 14–25 (model2)
 ]
 
-# ✅ MMIM Model definition with Swin-T backbone
-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 3 models trained on subsets of classes
-model1 = MMIM(num_classes=9)    # class 1–9
-model3 = MMIM(num_classes=4)    # class 10–13
-model2 = MMIM(num_classes=12)   # class 14–25
-
-model1.load_state_dict(torch.load("MMIM_best1.pth", map_location='cpu'))
-model3.load_state_dict(torch.load("MMIM_best3.pth", map_location='cpu'))
-model2.load_state_dict(torch.load("MMIM_best2.pth", map_location='cpu'))
-
-model1.eval()
-model2.eval()
-model3.eval()
-
-# ✅ Image preprocessing for Swin
+# ✅ Image transform for Swin
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -54,34 +23,26 @@ transform = transforms.Compose([
                          std=[0.5, 0.5, 0.5])
 ])
 
-# ✅ Prediction function
-def predict(image):
-    image_tensor = transform(image).unsqueeze(0)  # shape: [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]
+# ✅ Load models
+model1, model2, model3 = load_all_models()
 
-    # Combine predictions into one 25-class vector
-    combined = torch.cat([out1, out3, out2], dim=1)   # [1, 25]
-    pred_idx = combined.argmax(dim=1).item()
-    confidence = combined[0, pred_idx].item()
+# ✅ Gradio prediction function
+def predict(image):
+    label, confidence = predict_image(image, model1, model2, model3, transform, class_names)
 
     if confidence < 0.5:
-        return "Prediction uncertain – possibly unknown or low confidence"
+        return "⚠️ Prediction uncertain – possibly unknown or low confidence"
 
-    return f"🧠 Predicted: **{class_names[pred_idx]}** (Confidence: {confidence:.2f})"
+    return f"🧠 Predicted: **{label}** (Confidence: {confidence:.2f})"
 
-# ✅ Gradio UI
+# ✅ Gradio interface
 app = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload a weed image"),
     outputs=gr.Markdown(),
     title="🌿 Weed Classifier (25 Classes - MMIM)",
-    description="Upload an image to classify it into 1 of 25 weed classes using 3 separate MMIM models."
+    description="Upload an image to classify it into one of 25 weed categories using 3 Swin-MMIM models."
 )
 
-# ✅ Launch app
 if __name__ == "__main__":
     app.launch()