| 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} |
| |
| |
| image = image.convert("RGB") |
| |
| image = transform(image) |
| |
| image = image.unsqueeze(0).to(device) |
|
|
| |
| 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) |
