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

# Load model
model = resnet50(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, 5)
model.load_state_dict(torch.load("resnet50_dr.pth", map_location="cpu"))
model.eval()

class_names = ["No DR", "Mild", "Moderate", "Severe", "Proliferative DR"]
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

def predict(image):
    image = image.convert("RGB")
    img_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img_tensor)
        _, predicted = torch.max(outputs, 1)
        return class_names[predicted.item()]

gr.Interface(fn=predict, inputs="image", outputs="text").launch()