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()