Update app.py

app.py

@@ -7,9 +7,8 @@ from PIL import Image
 
 # 🔧 Model definition
 class MMIM(nn.Module):
-    def __init__(self, num_classes=12):
+    def __init__(self, num_classes=9):
         super(MMIM, self).__init__()
-        print("[INFO] Initializing MMIM model...")
         self.backbone = swin_t(weights='IMAGENET1K_V1')
         self.backbone.head = nn.Identity()
         self.classifier = nn.Sequential(
@@ -21,21 +20,16 @@ class MMIM(nn.Module):
 
     def forward(self, x):
         features = self.backbone(x)
-        print(f"[DEBUG] Backbone output shape: {features.shape}")
         return self.classifier(features)
 
 # ✅ Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"[INFO] Using device: {device}")
-
 model = MMIM(num_classes=12)
-print("[INFO] Loading model weights from MMIM_best2.pth...")
 model.load_state_dict(torch.load("MMIM_best2.pth", map_location=device))
 model.to(device)
 model.eval()
-print("[INFO] Model loaded and ready.")
 
-# ✅ Class names
+# ✅ Updated class names (match folder structure)
 class_names = [
     'Black grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common Wheat',
     'Fat Hen', 'Loose Silky-bent', 'Maize', 'Scentless Mayweed',
@@ -48,30 +42,31 @@ transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
-# 🔍 Prediction function
+# 🔍 Prediction function with negative detection
 def predict(img):
-    print("[INFO] Received image for prediction.")
     img = img.convert('RGB')
     img_tensor = transform(img).unsqueeze(0).to(device)
-    print(f"[DEBUG] Image tensor shape: {img_tensor.shape}")
 
     with torch.no_grad():
         outputs = model(img_tensor)
-        print(f"[DEBUG] Raw model outputs: {outputs}")
-        _, pred = torch.max(outputs, 1)
-        predicted_class = class_names[pred.item()]
-        print(f"[INFO] Predicted class: {predicted_class} (index {pred.item()})")
+        probs = torch.softmax(outputs, dim=1)
+        conf, pred = torch.max(probs, 1)
+
+    predicted_class = class_names[pred.item()]
+    confidence = conf.item() * 100
 
-    return predicted_class
+    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="label",
+    outputs="text",
     title="Weed Image Classifier",
-    description="Upload a weed image to predict its class"
+    description="Upload a weed image to predict its class. If the model detects a non-weed image, it will return 'Negative'."
 )
 
 interface.launch()
-