import torch
import torch.nn as nn
from torchvision import transforms, models
from PIL import Image
import gradio as gr

# === Load trained model ===
model = models.resnet18()
in_features = model.fc.in_features
model.fc = nn.Sequential(
    nn.Linear(in_features, 256),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(256, 2)
)
model.load_state_dict(torch.load("tumor_model.pth", map_location=torch.device("cpu")))
model.eval()

# === Transform (same as validation) ===
transform = transforms.Compose([
    transforms.Lambda(lambda x: x.convert('RGB')),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# === Prediction Function ===
def predict(image):
    image = Image.fromarray(image)
    input_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(input_tensor)
        _, pred = torch.max(outputs, 1)
        prob = torch.softmax(outputs, dim=1)[0][pred.item()].item()
        label = "Tumor: Yes" if pred.item() == 1 else "Tumor: No"
        return f"{label} ({prob * 100:.2f}%)"

# === Gradio Interface (No examples) ===
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Brain Scan"),
    outputs=gr.Label(label="Prediction"),
    title="🧠 Tumor Detection",
    description="Upload a brain MRI image to detect if a tumor is present (Yes or No)."
)

if __name__ == "__main__":
    interface.launch()