Update app.py

app.py

@@ -4,8 +4,6 @@ import torch.nn as nn
 from torchvision import transforms
 from torchvision.models import swin_t
 from PIL import Image
-import matplotlib.pyplot as plt
-import io
 
 # 🔧 Model definition
 class MMIM(nn.Module):
@@ -31,7 +29,7 @@ model.load_state_dict(torch.load("MMIM_best.pth", map_location=device))
 model.to(device)
 model.eval()
 
-# ✅ Class names
+# ✅ Updated class names (match folder structure)
 class_names = [
     "Chinee apple", "Lantana", "Negative", "Parkinsonia", "Parthenium",
     "Prickly acacia", "Rubber vine", "Siam weed", "Snake weed"
@@ -43,57 +41,32 @@ transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
-# 🔍 Prediction with confidence + top-3 chart
+# 🔍 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).cpu().squeeze()
+        probs = torch.softmax(outputs, dim=1)
+        conf, pred = torch.max(probs, 1)
 
-    conf, pred = torch.max(probs, 0)
     predicted_class = class_names[pred.item()]
     confidence = conf.item() * 100
 
-    # ✅ Main prediction message
     if predicted_class.lower() == "negative":
-        message = f"⚠️ Predicted as **Negative**\nConfidence: **{confidence:.2f}%**"
-    else:
-        message = f"✅ Predicted class: **{predicted_class}**\nConfidence: **{confidence:.2f}%**"
+        return f"⚠️ This image is predicted as Negative.\nConfidence: {confidence:.2f}%"
 
-    # 📊 Plot top 3 predictions
-    top_probs, top_idxs = torch.topk(probs, 3)
-    top_classes = [class_names[i] for i in top_idxs]
-    top_confidences = [p.item() * 100 for p in top_probs]
+    return f"✅ Predicted as a weed with class-{predicted_class}\nConfidence: {confidence:.2f}%"
 
-    fig, ax = plt.subplots(figsize=(6, 3))
-    ax.barh(top_classes[::-1], top_confidences[::-1], color='skyblue')
-    ax.set_xlim(0, 100)
-    ax.set_xlabel("Confidence (%)")
-    ax.set_title("Top 3 Predictions")
-    plt.tight_layout()
-
-    buf = io.BytesIO()
-    plt.savefig(buf, format="png")
-    plt.close(fig)
-    buf.seek(0)
-
-    from PIL import Image as PILImage
-    chart_img = PILImage.open(buf)
-
-    return message, chart_img
-
-# 🎨 Gradio UI
+# 🎨 Gradio Interface
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Textbox(label="Prediction with Confidence"),
-        gr.Image(type="pil", label="Top 3 Prediction Chart")
-    ],
-    title="Weed Image Classifier - Debug Mode",
-    description="🧪 Debug view enabled: see prediction confidence and top 3 predicted classes."
+    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()
+remove the confidence level