Tutorial

Automated Diagnosis of Retinal Diseases using Fundus Images and ResNet-50

Difficulty: Beginner Time: 12 min read

Introduction

Diabetic Retinopathy (DR) and Glaucoma are leading causes of blindness worldwide. Retinal fundus photography provides a non-invasive way to examine the back of the eye. In this tutorial, we will use PyTorch and a pre-trained ResNet-50 architecture to classify fundus images for the presence of retinal diseases, mimicking the workflow of a high-throughput ophthalmology screening system.

Architecture Overview


graph TD
    Fundus(Raw Fundus Image) --> Resize[Resize to 224x224]
    Resize --> Augment[Random Rotations & Flips]
    Augment --> Norm[Normalize RGB Channels]
    
    Norm --> ResNet(ResNet-50 Backbone - Frozen)
    ResNet --> Features(Feature Map 2048-dim)
    
    Features --> FC(New Fully Connected Layer - Trainable)
    FC --> Pred(5-Class Prediction: DR Stages)
    
    Pred -.-> |Interpretability| GradCAM(Grad-CAM Activation Map)
    GradCAM -.-> |Overlay| Heatmap(Lesion Heatmap on Original Image)

Prerequisites

  • Python 3.8+
  • PyTorch, Torchvision
  • Pillow (PIL) for image loading
  • A dataset of fundus images (e.g., the Kaggle Diabetic Retinopathy dataset or APTOS 2019)

Step 1: Install Dependencies

pip install torch torchvision pillow matplotlib opencv-python numpy
pip install grad-cam

Step 2: Data Preprocessing and Augmentation

Fundus images often have varying lighting conditions and artifacts. We use Torchvision transforms to resize the images to the standard ImageNet size (224x224), apply data augmentation (random flips and rotations) to prevent overfitting, and normalize the pixel values.

import torch
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

# Define transformations for training and validation
train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    # Standard ImageNet normalization
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 
])

val_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Assuming data is arranged in standard ImageFolder format (folders for each class)
# train_data = datasets.ImageFolder('./data/train', transform=train_transforms)
# val_data = datasets.ImageFolder('./data/val', transform=val_transforms)
# train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
# val_loader = DataLoader(val_data, batch_size=32, shuffle=False)

Step 3: Loading the Pre-trained ResNet-50

Instead of training a model from scratch, we use Transfer Learning. We load a ResNet-50 model pre-trained on ImageNet and replace the final classification layer to match the number of disease classes in our dataset (e.g., 5 stages of Diabetic Retinopathy).

import torchvision.models as models
import torch.nn as nn

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Load pre-trained ResNet-50
model = models.resnet50(pretrained=True)

# Freeze the early layers to retain generic feature extractors (optional but recommended for small datasets)
for param in model.parameters():
    param.requires_grad = False

# Replace the final fully connected layer
num_ftrs = model.fc.in_features
num_classes = 5 # 0: No DR, 1: Mild, 2: Moderate, 3: Severe, 4: Proliferative DR
model.fc = nn.Linear(num_ftrs, num_classes)

model = model.to(device)

Step 4: Training the Classification Head

Because we froze the earlier layers, we only need to train the newly initialized model.fc layer. We use Cross-Entropy Loss and the Adam optimizer.

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
# Only optimize the parameters that require gradients (the final FC layer)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    
    # Example training loop structure (assuming train_loader is defined)
    # for inputs, labels in train_loader:
    #     inputs, labels = inputs.to(device), labels.to(device)
    #     
    #     optimizer.zero_grad()
    #     outputs = model(inputs)
    #     loss = criterion(outputs, labels)
    #     loss.backward()
    #     optimizer.step()
    #     
    #     running_loss += loss.item()
        
    print(f"Epoch {epoch+1}/{num_epochs} - Loss: {running_loss:.4f}")

Step 5: Interpretability with Grad-CAM

In medical AI, "black box" models are often unacceptable. Clinicians need to know why a model diagnosed Severe DR. We use Gradient-weighted Class Activation Mapping (Grad-CAM) to generate a heatmap over the original fundus image, highlighting the regions (like microaneurysms or exudates) that the ResNet focused on to make its decision.

from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import numpy as np
import cv2

# Target the final convolutional layer of ResNet-50
target_layers = [model.layer4[-1]]

# Initialize Grad-CAM
cam = GradCAM(model=model, target_layers=target_layers, use_cuda=torch.cuda.is_available())

# Assume `input_tensor` is a single normalized image tensor of shape (1, 3, 224, 224)
# and `rgb_img` is the original image scaled to [0, 1] for visualization
# grayscale_cam = cam(input_tensor=input_tensor, targets=None)[0, :]
# 
# # Overlay the heatmap on the original image
# visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
# 
# import matplotlib.pyplot as plt
# plt.imshow(visualization)
# plt.title("Grad-CAM: Lesion Focus Area")
# plt.show()

Conclusion

Transfer learning with standard architectures like ResNet-50 allows medical AI developers to rapidly prototype highly accurate diagnostic models. By leveraging generic features learned from ImageNet, we can train powerful classifiers for retinal diseases even when medical image datasets are relatively small. Crucially, integrating interpretability tools like Grad-CAM bridges the trust gap, allowing ophthalmologists to verify that the model is detecting genuine pathology rather than image artifacts.