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

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

Project Overview
Dataset
Project Structure
Methodology
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:

Takes a retinal fundus image as input
Segments all lesions (Microaneurysms, Haemorrhages, Exudates, Optic Disc)
Predicts the DR severity grade (0-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:

Green Channel Extraction: The green channel of the RGB image provides the best contrast for retinal structures
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
Gaussian Blur: 3x3 kernel for noise reduction
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:

Runs the trained segmentation model to detect lesions and compute pixel-level coverage
Runs the trained grading model to predict DR severity (0-4)
Maps the predicted grade to a pre-authored clinical narrative written by domain experts
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:

python

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

json

{
  "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)

Class	Dice Score	IoU Score
Background	0.5089	0.3419
Microaneurysms	0.0026	0.0013
Haemorrhages	0.4420	0.3045
Hard Exudates	0.5047	0.3500
Soft Exudates	0.3889	0.2859
Optic Disc	0.9289	0.8689
Mean (excl. BG)	0.4534	0.3621

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%

Grade	Precision	Recall	F1-Score	Support
No DR (0)	0.6250	0.7353	0.6757	34
Mild DR (1)	0.0000	0.0000	0.0000	5
Moderate DR (2)	0.4773	0.6562	0.5526	32
Severe DR (3)	0.7059	0.6316	0.6667	19
Proliferative DR (4)	1.0000	0.1538	0.2667	13

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

Component	Specification
GPU	NVIDIA A100 SXM4 40GB
CPU	30 vCPUs
RAM	200 GB
Storage	500 GB SSD
Platform	Lambda Labs Cloud
Cost	~$1.99/hour
Total Training Time	~40 minutes
Total Cost	~$2.00

Software Environment

Package	Version	Purpose
Python	3.12	Programming language
PyTorch	2.6.0+cu124	Deep learning framework
CUDA	12.4	GPU acceleration
segmentation_models_pytorch	0.5.0	SegFormer architecture
albumentations	latest	Image augmentation
torchvision	0.21.0	EfficientNet, transforms
scikit-learn	1.9.0	Metrics, train/test split
OpenCV	4.13.0	Image I/O, preprocessing
matplotlib	latest	Plotting and visualization
pandas	latest	Data handling (CSV)
NumPy	2.4.4	Numerical operations
transformers	latest	VLM model loading
bitsandbytes	latest	4-bit quantization
timm	latest	Pretrained model zoo

Install All Dependencies

bash

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.

bash

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

bash

# 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

python

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

python

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/`)

File	Description
`run_pipeline.py`	Main script — runs the entire pipeline end-to-end. Contains all model definitions, training loops, evaluation, preprocessing, and report generation. ~1400 lines.
`generate_reports.py`	Standalone script to generate clinical reports using already-trained models. Loads both model weights, runs inference on test set, outputs JSON reports.
`setup_lambda.sh`	Shell script to set up a Lambda Labs GPU instance — installs all Python packages and creates required directories.

Results (`results/`)

File	Description
`segmentation/best_segformer.pth`	Trained SegFormer-B2 model weights (94 MB). Load with PyTorch.
`segmentation/training_history.png`	Plot showing loss, Dice, and IoU curves over 60 epochs.
`segmentation/per_class_metrics.png`	Bar chart comparing Dice and IoU for each lesion class.
`segmentation/test_metrics.csv`	CSV with exact Dice and IoU values per class.
`segmentation/predictions/*.png`	Color-coded predicted segmentation masks for all 27 test images.
`segmentation/visualizations/test_predictions_grid.png`	Side-by-side comparison of input, ground truth, and prediction for sample images.
`grading/best_efficientnet.pth`	Trained EfficientNet-B0 weights (21 MB).
`grading/grading_history.png`	Loss and accuracy curves over 30 epochs.
`grading/confusion_matrix.png`	5x5 confusion matrix showing classification errors.
`grading/classification_report.txt`	Full precision/recall/F1 report from scikit-learn.
`grading/test_predictions.csv`	Per-image predictions with true vs predicted grades.
`vlm_reports/all_reports.json`	27 structured clinical reports with lesion analysis.

References

IDRiD Dataset: Porwal, P., et al. "IDRiD: Diabetic Retinopathy – Segmentation and Grading Challenge." Medical Image Analysis, 2020.
SegFormer: Xie, E., et al. "SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers." NeurIPS 2021.
EfficientNet: Tan, M. & Le, Q. "EfficientNet: Rethinking Model Scaling for CNNs." ICML 2019.
Focal Loss: Lin, T.Y., et al. "Focal Loss for Dense Object Detection." ICCV 2017.
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.

raiyanz/diabetic-retinopathy-analysis

Diabetic Retinopathy Analysis Pipeline

Table of Contents

Project Overview

Dataset

Dataset Statistics

Segmentation Classes (6 classes)

Disease Grading Classes (5 grades)

Dataset Folder Structure

Project Structure

Methodology

Preprocessing Pipeline

Stage 1: Multi-Class Lesion Segmentation

Model Architecture

Why SegFormer?

Training Configuration

Loss Function: Combined Dice + Focal Loss

Class Weights (to handle extreme imbalance)

Data Augmentation (Critical for 54 images)

Multi-Class Mask Construction

Stage 2: Disease Severity Grading

Model Architecture

Why EfficientNet?

Training Configuration

Grading Data Augmentation

Stage 3: Clinical Report Generation (VLM)

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

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

Alternative: Google Gemini Vision API

Report Content

Example Report Output

Results

Segmentation Results (Test Set — 27 images)

Disease Grading Results (Test Set — 103 images)

VLM Report Generation Results

Environment & Setup

Hardware Used for Training

Software Environment

Install All Dependencies

How to Run

Option 1: Run the Full Pipeline (GPU Required)

Option 2: Run on Lambda Labs (Recommended)

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

How to Use the Trained Models

Load Segmentation Model

Load Grading Model

File Descriptions

Training Code (training/)

Results (results/)

References

License

Contact

raiyanz/
diabetic-retinopathy-analysis

Training Code (`training/`)

Results (`results/`)