diabetic-retinopathy-analysis

Diabetic Retinopathy Analysis Pipeline

A complete end-to-end deep learning pipeline for Diabetic Retinopathy (DR) analysis using retinal fundus images. This project performs three tasks:

1. Multi-class Retinal Lesion Segmentation (SegFormer-B2)
2. Disease Severity Grading (EfficientNet-B0)
3. Automated Clinical Report Generation (VLM-based)

---

Table of Contents

• Project Overview
• Dataset
• Project Structure
• Methodology
  • Stage 1: Segmentation
  • Stage 2: Disease Grading
  • Stage 3: VLM Report Generation
• Results
• Environment & Setup
• How to Run
• How to Use the Trained Models
• File Descriptions
• References

---

Project Overview

Diabetic Retinopathy is a leading cause of blindness worldwide. Early detection through retinal imaging can prevent vision loss. This project builds an automated pipeline that:

1. Takes a retinal fundus image as input
2. Segments all lesions (Microaneurysms, Haemorrhages, Exudates, Optic Disc)
3. Predicts the DR severity grade (0-4)
4. Generates a clinical report explaining the findings

The entire pipeline was trained on the IDRiD (Indian Diabetic Retinopathy Image Dataset) using a single NVIDIA A100 40GB GPU on Lambda Labs cloud.

---

Dataset

Name: IDRiD (Indian Diabetic Retinopathy Image Dataset)

Source: IEEE DataPort / Grand Challenge

The dataset contains high-resolution retinal fundus photographs with pixel-level annotations for lesions and image-level disease grades.

Dataset Statistics

Segmentation Classes (6 classes)

Disease Grading Classes (5 grades)

Dataset Folder Structure

dataset/
├── A. Segmentation/
│   ├── 1. Original Images/
│   │   ├── a. Training Set/          # 54 RGB fundus images (.jpg)
│   │   └── b. Testing Set/           # 27 RGB fundus images (.jpg)
│   └── 2. All Segmentation Groundtruths/
│       ├── a. Training Set/
│       │   ├── 1. Microaneurysms/    # Binary masks (_MA.tif)
│       │   ├── 2. Haemorrhages/      # Binary masks (_HE.tif)
│       │   ├── 3. Hard Exudates/     # Binary masks (_EX.tif)
│       │   ├── 4. Soft Exudates/     # Binary masks (_SE.tif)
│       │   └── 5. Optic Disc/        # Binary masks (_OD.tif)
│       └── b. Testing Set/
│           └── (same 5 folders)
└── B. Disease Grading/
    ├── 1. Original Images/
    │   ├── a. Training Set/          # 413 fundus images
    │   └── b. Testing Set/           # 103 fundus images
    └── 2. Groundtruths/
        ├── a. IDRiD_Disease Grading_Training Labels.csv
        └── b. IDRiD_Disease Grading_Testing Labels.csv

---

Project Structure

diabetic-retinopathy-analysis/
│
├── README.md                          # This file (complete project documentation)
│
├── dataset/                           # IDRiD dataset (images + ground truths)
│   ├── A. Segmentation/              # Segmentation task data
│   └── B. Disease Grading/           # Grading task data
│
├── training/                          # All training code
│   ├── run_pipeline.py               # Main pipeline script (runs everything)
│   ├── generate_reports.py           # VLM clinical report generator
│   └── setup_lambda.sh              # GPU server setup script
│
└── results/                           # All outputs from training
    ├── segmentation/
    │   ├── best_segformer.pth         # Trained segmentation model weights
    │   ├── training_history.png       # Loss and Dice curves over epochs
    │   ├── per_class_metrics.png      # Bar chart of Dice/IoU per class
    │   ├── test_metrics.csv           # Numerical metrics
    │   ├── predictions/               # Predicted masks for all 27 test images
    │   └── visualizations/            # GT vs Prediction comparison grid
    ├── grading/
    │   ├── best_efficientnet.pth      # Trained grading model weights
    │   ├── grading_history.png        # Loss and accuracy curves
    │   ├── confusion_matrix.png       # 5x5 confusion matrix
    │   ├── classification_report.txt  # Precision, Recall, F1 per class
    │   └── test_predictions.csv       # Per-image predictions
    └── vlm_reports/
        └── all_reports.json           # Clinical reports for 27 test images

---

Methodology

Preprocessing Pipeline

Every fundus image goes through the following preprocessing before being fed to any model:

1. Green Channel Extraction: The green channel of the RGB image provides the best contrast for retinal structures
2. Multi-Scale CLAHE: Contrast Limited Adaptive Histogram Equalization applied at 3 tile grid sizes (4x4, 8x8, 16x16) and blended (25%, 50%, 25%) for uniform enhancement across varying illumination
3. Gaussian Blur: 3x3 kernel for noise reduction
4. Resize: 512x512 pixels for segmentation, 224x224 for grading

---

Stage 1: Multi-Class Lesion Segmentation

Model Architecture

• Model: SegFormer-B2 (from segmentation_models_pytorch library)
• Encoder: MiT-B2 (Mix Transformer) pretrained on ImageNet
• Input: 1-channel grayscale (repeated to 3 channels for compatibility with pretrained encoder)
• Output: 6-class pixel-wise segmentation map (512x512)
• Total Parameters: 24.7 Million

Why SegFormer?

SegFormer combines a hierarchical Transformer encoder with a lightweight MLP decoder. It captures both local details (important for tiny microaneurysms) and global context (important for understanding overall retinal structure). It outperforms U-Net on medical image segmentation with fewer parameters.

Training Configuration

Loss Function: Combined Dice + Focal Loss

Total Loss = 0.5 * Dice Loss + 0.5 * Focal Loss

• Dice Loss: Directly optimizes the overlap between prediction and ground truth. Handles class imbalance well.
• Focal Loss (gamma=2.0): Down-weights easy-to-classify pixels, focuses training on hard pixels (lesion boundaries).

Class Weights (to handle extreme imbalance)

Data Augmentation (Critical for 54 images)

Since we only have 54 training images, heavy augmentation is essential:

• HorizontalFlip (p=0.5)
• VerticalFlip (p=0.3)
• RandomRotate90 (p=0.5)
• ShiftScaleRotate (shift=0.05, scale=0.1, rotate=20deg, p=0.5)
• GridDistortion / ElasticTransform (p=0.3)
• RandomBrightnessContrast (p=0.4)
• GaussNoise (var=5-30, p=0.3)

Multi-Class Mask Construction

Individual binary masks from the dataset are combined into a single multi-class mask with priority ordering (higher class ID overwrites lower if overlapping):

Background (0) → MA (1) → HE (2) → EX (3) → SE (4) → OD (5)

---

Stage 2: Disease Severity Grading

Model Architecture

• Model: EfficientNet-B0 (pretrained on ImageNet)
• Input: 224x224 RGB fundus image
• Output: 5-class probability (Grades 0-4)
• Final Layer: Linear(1280, 5)

Why EfficientNet?

EfficientNet uses compound scaling (depth, width, resolution) to achieve high accuracy with fewer parameters. B0 is the smallest variant — ideal for a dataset of 413 images.

Training Configuration

Grading Data Augmentation

• RandomHorizontalFlip
• RandomVerticalFlip
• RandomRotation (15 degrees)
• ColorJitter (brightness=0.2, contrast=0.2)

---

Stage 3: Clinical Report Generation (VLM)

This stage uses a Vision-Language Model (VLM) approach to generate structured clinical reports from the model outputs. Two approaches were explored:

Approach 1: Template-Based VLM (Used in Final Pipeline)

The production pipeline uses a template-based deterministic approach that combines segmentation and grading outputs to generate clinically accurate reports without requiring an external LLM API:

1. Runs the trained segmentation model to detect lesions and compute pixel-level coverage
2. Runs the trained grading model to predict DR severity (0-4)
3. Maps the predicted grade to a pre-authored clinical narrative written by domain experts
4. Appends lesion-specific findings (name + coverage %) with clinical interpretation

This approach was chosen because:

• Deterministic output — same input always produces the same report (important for medical AI)
• No hallucination risk — all text is pre-validated by clinical templates
• No API dependency — runs fully offline on any hardware
• Fast inference — no LLM generation latency

Approach 2: LLaVA Fine-tuning (Experimental, Not Deployed)

An attempt was made to fine-tune LLaVA (Large Language and Vision Assistant) on the IDRiD dataset for open-ended report generation:

• Training data was formatted as conversation-style JSON (human asks for diagnosis, model provides report)
• The 03_dr_grading_vlm.py notebook prepares the LLaVA-compatible training JSON from the grading labels
• Due to compute constraints (LLaVA-7B requires >24GB VRAM for fine-tuning), this was explored in Google Colab but not deployed in the final pipeline
• The formatted training JSONs (idrid_vlm_train.json, idrid_vlm_val.json) are available for future fine-tuning

Alternative: Google Gemini Vision API

For richer, context-aware reports without fine-tuning, you can use the Gemini Vision API as a drop-in replacement:

import google.generativeai as genai
from PIL import Image

genai.configure(api_key="YOUR_GEMINI_API_KEY")
model = genai.GenerativeModel("gemini-2.0-flash")

img = Image.open("fundus_image.jpg")

prompt = """You are an ophthalmologist. Analyze this retinal fundus image.
The segmentation model detected: Microaneurysms (58.08%), Hard Exudates (6.30%).
The grading model predicted: Severe DR (Grade 3).

Provide a structured clinical report with:
1. Overall assessment
2. Lesion analysis
3. Clinical significance
4. Recommended follow-up"""

response = model.generate_content([prompt, img])
print(response.text)

Gemini Vision API requirements:

• pip install google-generativeai
• A Google AI Studio API key (free tier: 15 RPM / 1M tokens per day)
• Internet connection for API calls

This gives more natural, detailed reports but introduces non-determinism and API dependency.

Report Content

Each generated report includes:

• Overall DR severity assessment
• Description of each detected lesion type
• Clinical significance of the findings
• Coverage percentage of each lesion
• Recommended follow-up actions

Example Report Output

{
  "image_id": "IDRiD_55",
  "dr_grade": 3,
  "grade_name": "Severe DR",
  "lesions_detected": [
    {"name": "Microaneurysms", "percentage": 58.08},
    {"name": "Haemorrhages", "percentage": 0.67},
    {"name": "Hard Exudates", "percentage": 6.30},
    {"name": "Optic Disc", "percentage": 1.53}
  ],
  "clinical_report": "Severe non-proliferative diabetic retinopathy is detected. 
   Extensive retinal lesions are observed with significant hemorrhages, cotton wool 
   spots, and venous abnormalities. Urgent referral to a retinal specialist is 
   recommended. Microaneurysms detected (58.08% coverage), indicating focal 
   capillary wall weakness. Retinal hemorrhages observed (0.67% coverage), 
   suggesting vessel wall rupture. Hard exudates present (6.30% coverage), 
   indicating lipid leakage from damaged vessels."
}

---

Results

Segmentation Results (Test Set — 27 images)

Key Observations:

• Optic Disc achieves excellent segmentation (Dice 0.93) due to its large, consistent shape
• Hard Exudates perform well (Dice 0.50) as they have bright, distinct appearance
• Microaneurysms are nearly impossible to segment (Dice 0.003) because they are only 2-5 pixels in diameter
• With only 54 training images, these results are reasonable for the IDRiD benchmark

Disease Grading Results (Test Set — 103 images)

Overall Test Accuracy: 58.25%

Key Observations:

• Severe DR has the best F1-score (0.67) with good precision and recall
• No DR is well-classified (F1 0.68) since healthy retinas are distinct
• Mild DR has 0 samples correctly predicted (only 5 test samples — too few to learn)
• Class imbalance is the main challenge (20 mild samples vs 136 moderate in training)

VLM Report Generation Results

• 27 clinical reports generated for all test images
• Each report includes lesion-level analysis with coverage percentages
• Reports follow standard ophthalmological format
• Available in results/vlm_reports/all_reports.json

---

Environment & Setup

Hardware Used for Training

Software Environment

Install All Dependencies

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install segmentation_models_pytorch albumentations opencv-python-headless
pip install scikit-learn matplotlib pandas pillow timm tqdm
pip install transformers accelerate bitsandbytes

---

How to Run

Option 1: Run the Full Pipeline (GPU Required)

This trains both models from scratch and generates reports.

# 1. Clone this repo
git clone https://github.com/YOUR_USERNAME/diabetic-retinopathy-analysis.git
cd diabetic-retinopathy-analysis

# 2. Install dependencies
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install segmentation_models_pytorch albumentations opencv-python-headless
pip install scikit-learn matplotlib pandas pillow timm tqdm

# 3. Run the pipeline
python training/run_pipeline.py

The script will:

• Verify dataset paths
• Preprocess all images (Green channel + CLAHE)
• Build multi-class segmentation masks
• Train SegFormer-B2 for 60 epochs
• Evaluate segmentation on test set
• Train EfficientNet-B0 for 30 epochs
• Evaluate grading on test set
• Generate clinical reports

All results are saved to ~/results/.

Option 2: Run on Lambda Labs (Recommended)

# 1. Launch 1x A100 40GB instance on Lambda Labs
# 2. SSH into the instance
ssh ubuntu@YOUR_IP

# 3. Upload dataset
scp -r dataset/ ubuntu@YOUR_IP:~/dataset/

# 4. Upload training code
scp training/run_pipeline.py ubuntu@YOUR_IP:~/

# 5. Setup and run
bash setup_lambda.sh
python run_pipeline.py

# 6. Download results
scp -r ubuntu@YOUR_IP:~/results/ ./results/

Option 3: Generate Reports Only (Using Pre-trained Models)

If you already have the trained model weights:

python training/generate_reports.py

This loads best_segformer.pth and best_efficientnet.pth, runs inference on test images, and generates clinical reports.

---

How to Use the Trained Models

Load Segmentation Model

import torch
import torch.nn as nn
import segmentation_models_pytorch as smp

class SegFormerModel(nn.Module):
    def __init__(self, num_classes=6):
        super().__init__()
        self.model = smp.create_model(
            arch="segformer",
            encoder_name="mit_b2",
            encoder_weights=None,  # We load our own weights
            in_channels=3,
            classes=num_classes,
        )

    def forward(self, x):
        # Input: (B, 1, 512, 512) grayscale
        x = x.repeat(1, 3, 1, 1)  # Convert to 3-channel
        return self.model(x)

# Load model
model = SegFormerModel(num_classes=6)
model.load_state_dict(torch.load("results/segmentation/best_segformer.pth", map_location="cpu"))
model.eval()

# Inference
import cv2
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2

transform = A.Compose([
    A.Resize(512, 512),
    A.Normalize(mean=(0.5,), std=(0.25,)),
    ToTensorV2(),
])

# Preprocess: extract green channel + CLAHE
img_bgr = cv2.imread("your_fundus_image.jpg")
green = img_bgr[:, :, 1]
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
enhanced = clahe.apply(green)

# Transform and predict
tensor = transform(image=enhanced)["image"].unsqueeze(0).float()
with torch.no_grad():
    logits = model(tensor)
    prediction = torch.argmax(logits, dim=1).squeeze().numpy()

# prediction is a 512x512 array with values 0-5
# 0=Background, 1=MA, 2=HE, 3=EX, 4=SE, 5=OD

Load Grading Model

from torchvision import models, transforms
from PIL import Image

# Load model
model = models.efficientnet_b0(weights=None)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 5)
model.load_state_dict(torch.load("results/grading/best_efficientnet.pth", map_location="cpu"))
model.eval()

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

img = Image.open("your_fundus_image.jpg").convert("RGB")
tensor = transform(img).unsqueeze(0)

with torch.no_grad():
    output = model(tensor)
    grade = torch.argmax(output, dim=1).item()

# grade is 0-4 (No DR, Mild, Moderate, Severe, Proliferative)
print(f"DR Grade: {grade}")

---

File Descriptions

Training Code (training/)

Results (results/)

---

References

1. IDRiD Dataset: Porwal, P., et al. "IDRiD: Diabetic Retinopathy – Segmentation and Grading Challenge." Medical Image Analysis, 2020.
2. SegFormer: Xie, E., et al. "SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers." NeurIPS 2021.
3. EfficientNet: Tan, M. & Le, Q. "EfficientNet: Rethinking Model Scaling for CNNs." ICML 2019.
4. Focal Loss: Lin, T.Y., et al. "Focal Loss for Dense Object Detection." ICCV 2017.
5. CLAHE: Pizer, S.M., et al. "Adaptive histogram equalization and its variations." Computer Vision, Graphics, and Image Processing, 1987.

---

License

This project uses the IDRiD dataset which is released under the CC BY 4.0 license.

---

Contact

For questions about this project, feel free to open an issue on this repository.