Update app.py

app.py

@@ -35,7 +35,7 @@ model.eval()
 print("[INFO] Model loaded successfully.")
 
 # ✅ Class names
-class_names = ["Broadleaf", "Grass", "Soil", "Soybean"]
+class_names = []
 
 # 🔁 Transform
 transform = transforms.Compose([
@@ -73,3 +73,75 @@ interface = gr.Interface(
 )
 
 interface.launch()
+
+
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from torchvision.models import swin_t
+from PIL import Image
+
+# 🔧 Model definition
+class MMIM(nn.Module):
+    def __init__(self, num_classes=9):
+        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):
+        features = self.backbone(x)
+        return self.classifier(features)
+
+# ✅ Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = MMIM(num_classes=4)
+model.load_state_dict(torch.load("MMIM_best3.pth", map_location=device))
+model.to(device)
+model.eval()
+
+# ✅ Updated class names (match folder structure)
+class_names = [
+  "Broadleaf", "Grass", "Soil", "Soybean"
+]
+
+# 🔁 Image transform
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor()
+])
+
+# 🔍 Prediction function with negative detection
+def predict(img):
+    img = img.convert('RGB')
+    img_tensor = transform(img).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probs = torch.softmax(outputs, dim=1)
+        conf, pred = torch.max(probs, 1)
+
+    predicted_class = class_names[pred.item()]
+    confidence = conf.item() * 100
+
+    if predicted_class.lower() == "negative":
+        return f"⚠️ This image is predicted as Negative.\nConfidence: {confidence:.2f}%"
+
+    return f"✅ Predicted as a weed with class-{predicted_class}\nConfidence: {confidence:.2f}%"
+
+# 🎨 Gradio Interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Weed Image Classifier",
+    description="Upload a weed image to predict its class. If the model detects a non-weed image, it will return 'Negative'."
+)
+
+interface.launch()