import torch
import torch.nn as nn
from torchvision import transforms,models
from PIL import Image
import gradio as gr
class_name=["glioma","meningioma","notumor","pituitary"]
# Load the model
model = models.resnet50(pretrained=False)
model.fc = nn.Linear(model.fc.in_features, len(class_name))
#model.classifier[6] = nn.Linear(model.classifier[6].in_features, len(class_name))
checkpoint = torch.load('brain_tumor_model.pth', map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Define the transforms
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Function to make predictions
def predict_image(img):
    
    img = transform(img).unsqueeze(0)  # Preprocess the image
    with torch.no_grad():
        outputs = model(img)
        _, predicted = torch.max(outputs, 1)
    return class_name[predicted.item()]

# Gradio Interface
gr.Interface(fn=predict_image, inputs=gr.Image(type="pil"), outputs="label").launch()