This project performs semantic segmentation of blood vessels in retinal fundus images using a custom U-Net architecture. The model is trained on the DRIVE dataset to assist in early diagnosis of diabetic retinopathy and other retinal diseases.
- Custom U-Net implementation in PyTorch
- Combined BCE + Dice loss for improved performance
- Data augmentation with Albumentations
- Streamlit-powered interactive demo
- Evaluation metrics: Dice Score ~0.76, IoU ~0.62
Place the DRIVE dataset inside the dataset/
directory in the following format:
dataset/
├── training/
│ ├── images/
│ └── masks/
├── test/
│ ├── images/
│ └── masks/
git clone https://github.com/Suyog-16/retinal-vessel-segmentation.git
cd retinal-vessel-segmentation
conda env create -f environment.yml
conda activate retinal-vessel-seg
Requirements:
- Python 3.10+
- PyTorch >= 1.9
- Albumentations
- Streamlit
- Numpy, matplotlib, etc.
python src/train.py --batch_size 2 --epochs 50 --lr 0.0001
CLI Arguments:
Argument | Description | Default |
---|---|---|
--batch_size |
Batch size for training | 2 |
--epochs |
Number of training epochs | 50 |
--lr |
Learning rate for optimizer | 0.0001 |
Model checkpoints are saved in models/
. Best checkpoint: models/best_unet.pth
.
python src/inference.py --img_path img/example.png --model_path models/best_unet.pth
This displays the segmentation output using matplotlib. You can also save the output by modifying the script.
streamlit run app/app.py
Use the web interface to upload images and visualize segmented blood vessels in real time.
Metric | Score on DRIVE Test Set |
---|---|
Dice Score | ~0.76 |
IoU Score | ~0.62 |
Input Image | Ground Truth | Model Prediction |
---|---|---|
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
This project uses a custom U-Net variant with the following key improvements over the original architecture:
- Each block is Conv → BatchNorm → ReLU (for better convergence)
- Modular
DoubleConv
blocks used for cleaner design - Standard encoder–decoder with skip connections
- Output is a binary segmentation map (1-channel)
- Trained with a compound loss: BCE + Dice Loss
This setup offers improved training stability and better performance on small datasets like DRIVE.
- U-Net: Convolutional Networks for Biomedical Image Segmentation
- Albumentations: A fast and flexible image augmentation library
- DRIVE dataset
This project is licensed under the MIT License.