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

device = torch.device("cpu")
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features,2)

model = model.to(device)

model.load_state_dict(
    torch.load(
        "resnet18_retinopathy.pth",
        map_location=device))

model.eval()


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

class_names = ["no DR", "DR present"]

def predict(image):
    print("Received image:", type(image))
    if image is None:
        return {"no DR": 0.0, "DR present": 0.0}
        
    # Make sure image is RGB
    image = image.convert("RGB")
    # Apply same preprocessing used during testing
    image = transform(image)
    # Add batch dimension
    image = image.unsqueeze(0).to(device)

    # Make prediction
    with torch.no_grad():
        output = model(image)
        probabilities = torch.softmax(output, dim=1)[0]

    result = {
        "no DR": float(probabilities[0].item()),
        "DR present": float(probabilities[1].item())}
        
    print("Prediction:", result)
    return result

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Diabetic Retinopathy Classifier",
    description="Upload a retinal image to obtain a prediction ")

demo.launch(
    debug=True,
    show_error=True,
    ssr_mode=False)