DIAGNijmegen / drive-vessels-unet

Retinal vessel segmentation on DRIVE

This repository will detail how you can train and deploy a U-Net for retinal vessel segmentation on the DRIVE dataset using Docker.

💡 To use the pre-trained model and wrap it in a Docker container, follow the instructions provided in the documentation for grand-challenge.org.

Requirements (for training)

• monai
• SimpleITK
• numpy
• scipy
• skimage
• torch
• torchvision

Training

For training the algorithm, first download the DRIVE dataset and place the files under data/. You'll need to sign up for the Challenge through grand-challenge.org to access the link to the dataset.

Start training your algorithm by executing the following:

python train.py

Pre-trained weights

Pre-trained weights are available in this repository under the name best_metric_model_segmentation2d_dict.pth.

Inference

Run inference.py to take a test image and plot the prediction along with the input image.