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

Split	Images	Purpose
Segmentation Training	54	Train segmentation model
Segmentation Testing	27	Evaluate segmentation
Grading Training	413	Train grading model
Grading Testing	103	Evaluate grading

Segmentation Classes (6 classes)

Class ID	Class Name	Description
0	Background	Normal retinal tissue
1	Microaneurysms (MA)	Tiny bulges in blood vessel walls (earliest DR sign)
2	Haemorrhages (HE)	Blood leakage from damaged retinal vessels
3	Hard Exudates (EX)	Yellow lipid deposits from leaking vessels
4	Soft Exudates (SE)	Cotton wool spots (nerve fiber layer infarcts)
5	Optic Disc (OD)	Optic nerve head (anatomical landmark)

Disease Grading Classes (5 grades)

Grade	Name	Description
0	No DR	Healthy retina, no signs of disease
1	Mild NPDR	Few microaneurysms only
2	Moderate NPDR	Multiple hemorrhages, exudates present
3	Severe NPDR	Extensive lesions, venous abnormalities
4	Proliferative DR	Neovascularization, high risk of blindness

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

Parameter	Value
Image Size	512 x 512
Batch Size	8
Epochs	60
Optimizer	AdamW (weight_decay=1e-4)
Learning Rate	3e-4
LR Scheduler	CosineAnnealingLR (eta_min=1e-6)
Mixed Precision	Yes (AMP with GradScaler)
Train/Val Split	80/20 (43 train, 11 val)
Test Set	27 images (separate, provided by IDRiD)

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)

Class	Weight	Reason
Background	0.02	Dominates 95%+ of pixels
Microaneurysms	4.0	Extremely tiny, <0.1% of pixels
Haemorrhages	3.5	Small lesions
Hard Exudates	2.0	Moderate size
Soft Exudates	3.0	Rare
Optic Disc	0.5	Large, easy to segment

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

Parameter	Value
Image Size	224 x 224
Batch Size	16
Epochs	30
Optimizer	Adam (lr=1e-4)
LR Scheduler	StepLR (step=10, gamma=0.5)
Loss	CrossEntropyLoss
Train/Val Split	80/20 stratified (330 train, 83 val)
Test Set	103 images (separate)

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...