The Little W-Net that Could

You have reached the official repository for our work on retinal vessel segmentation with minimalistic models. The above picture represents a WNet architecture, which contains roughly around 70k parameters and closely matches (or outperforms) other more complicated techniques. For more details about our work, you can check the related paper:

The Little W-Net That Could: State-of-the-Art Retinal Vessel Segmentation with Minimalistic Models
Adrian Galdran, André Anjos, Jose Dolz, Hadi Chakor, Hervé Lombaert, Ismail Ben Ayed
https://arxiv.org/abs/2009.01907, Sep. 2020

We would appreciate if you could cite our work if it is useful for you :)

Note: If you are just looking for our results, you can directly download them at this link.

Please find below a table of contents describing what you can find in this repository:

Table of Contents

1. Dependencies and getting the data ready
2. Training a W-Net for vessel segmentation
3. Generating segmentations
4. Computing Performance
5. Cross-Dataset Experiments
6. Training with pseudo-labels and computing performance
7. Evaluating your own model
8. Training a W-Net for Artery/Vein segmentation
9. Generating Artery/Vein segmentations
10. Generating vessel and A/V segmentations on your own data

1. Dependencies and getting the data ready

First things first, clone this repo somewhere in your computer:

git clone https://github.com/agaldran/lwnet.git .

For full reproducibility, you should use the configuration specified in the requirements.txt file. If you are using conda, you can install dependencies in one line, just run on a terminal:

conda create --name lwnet --file environment.txt
conda activate lwnet

We have made an effort to automate the data download and preparation so that everything is as reproducible as possible. Out of the ten datasets we use in the paper, seven of them are public, and you can get them just running:

python get_public_data.py

This will populate the data directory with the seven sub-folders. If everything goes right, each sub-folder in data is named as the corresponding dataset, and contains at least:

• Three folders called images, mask, manual
• A csv file called test_all.csv

If the dataset is used in our work for training a vessel segmentation model (DRIVE, CHASE-DB, and HRF), you will also find:

• Three csv files called train.csv, val.csv, test.csv

If the dataset also has Artery/Vein annotations, you will also see:

• A folder called manual_av
• A csv file called test_all_av.csv

If the dataset is used in our work for training an A/V models (DRIVE and HRF), you will also find:

• Three csv files called train_av.csv, val_av.csv, test_av.csv

Note: The DRIVE dataset will also contain a folder called ZoneB_manual, which is used to evaluate A/V performance around the optic disc. The HRF dataset will also contain folders called images_resized, manual_resized, mask_resized. These are used only for training.

Note: The LES-AV dataset is still public but it now needs to be downloaded manually, please see the comments in get_public_data.py Line 400 forward for details.

2. Training a W-Net for vessel segmentation

Train a model on a given dataset. You also need to supply the path to save the model. Note that the training defaults to using the CPU, which is feasible due to the small size of our models. To reproduce our results in table 2 of our paper, you need to run:

python train_cyclical.py --csv_train data/DRIVE/train.csv --cycle_lens 20/50
                         --model_name wnet --save_path wnet_drive --device cuda:0
python train_cyclical.py --csv_train data/CHASEDB/train.csv --cycle_lens 40/50
                         --model_name wnet --save_path wnet_chasedb --device cuda:0
python train_cyclical.py --csv_train data/HRF/train.csv --cycle_lens 30/50
                         --model_name wnet --save_path wnet_hrf_1024
                         --im_size 1024 --batch_size 2 --grad_acc_steps 1 --device cuda:0

This will store the model weights in experiments/wnet_drive, experiments/wnet_chasedb, experiments/wnet_hrf respectively.

The parameter cycle_lens specifies the length of the training, and it is adjusted depending on the amount of images in the training set. For instance, in the DRIVE case, --cycle_lens 20/50 implies that we train for 20 cycles, each cycle running for 50 epochs. As CHASE-DB has less training images than DRIVE (8 vs 16), we double the number of cycles in that case.

Note that we use a batch_size of 4 by default, and that we train on HRF with an image size of 1024x1024. In order to train on a single GPU, we use gradient accumulation in that case.

3. Generating segmentations

Once the model is trained, you can produce the corresponding segmentations calling generate_results.py and specifying which dataset should be used:

python generate_results.py --config_file experiments/wnet_drive/config.cfg
                           --dataset DRIVE --device cuda:0
python generate_results.py --config_file experiments/wnet_chasedb/config.cfg
                           --dataset CHASEDB --device cuda:0
python generate_results.py --config_file experiments/wnet_hrf_1024/config.cfg
                           --dataset HRF --im_size 1024 --device cuda:0

The above stores the predictions for those datasets in results/DRIVE/experiments/wnet_drive, results/CHASEDB/experiments/wnet_chasedb, and results/HRF/experiments/wnet_hrf_1024 respectively.

4. Computing Performance

We call analyze_results.py to compute performance. It is important to specify what was the training and what is the test set here. For that, you pass the path to the train/test predictions, and the name of the train/test datasets:

python analyze_results.py --path_train_preds results/DRIVE/experiments/wnet_drive
                          --path_test_preds results/DRIVE/experiments/wnet_drive
                          --train_dataset DRIVE --test_dataset DRIVE

python analyze_results.py --path_train_preds results/CHASEDB/experiments/wnet_chasedb
                          --path_test_preds results/CHASEDB/experiments/wnet_chasedb
                          --train_dataset CHASEDB --test_dataset CHASEDB

python analyze_results.py --path_train_preds results/HRF/experiments/wnet_hrf_1024
                          --path_test_preds results/HRF/experiments/wnet_hrf_1024
                          --train_dataset HRF --test_dataset HRF

The code uses the csv files in each dataset folder to check which images should be used for running an AUC analysis in the training set and finding an optimal binarizing threshold to be used in the test set images.

5. Cross-Dataset Experiments

When a model has been trained on dataset A (say, DRIVE) and we want to test it on dataset B (say, CHASE-DB), we first generate segmentations on both datasets:

python generate_results.py --config_file experiments/wnet_drive/config.cfg
                           --dataset DRIVE  --device cuda:0
python generate_results.py --config_file experiments/wnet_drive/config.cfg
                           --dataset CHASEDB  --device cuda:0

and then we compute performance:

python analyze_results.py --path_train_preds results/DRIVE/experiments/wnet_drive
                          --path_test_preds results/CHASEDB/experiments/wnet_drive
                          --train_dataset DRIVE --test_dataset CHASEDB

6. Training with pseudo-labels and computing performance

1. Train a model on a source dataset (DRIVE); this will store the model in experiments/wnet_drive

python train_cyclical.py --csv_train data/DRIVE/train.csv --cycle_lens 20/50
                         --model_name wnet --save_path wnet_drive
                         --device cuda:0

2. Generate predictions on target dataset (CHASEDB) with this model; this will store predictions at results/CHASEDB/experiments/wnet_drive

python generate_results.py --config_file experiments/wnet_drive/config.cfg
                           --dataset CHASEDB --device cuda:0

3. Train a model on DRIVE manual segmentations plus CHASEDB pseudo-segmentations for one cycle of 10 epochs with a lower learning rate, starting from the weights of the model trained on DRIVE. Note that in this case we use the AUC on the training set as checkpointing criterion. This training is slower because of the AUC computation on a large set of images at the end of each cycle. In this case, we save the new model in a folder called wnet_drive_chasedb_pl:

python train_cyclical.py --save_path wnet_drive_chasedb_pl
                         --checkpoint_folder experiments/wnet_drive
                         --csv_test data/CHASEDB/test_all.csv
                         --path_test_preds results/CHASEDB/experiments/wnet_drive
                         --max_lr 0.0001 --cycle_lens 10/1 --metric tr_auc
                         --device cuda:0

3. Generate predictions with this new model on source dataset DRIVE:

python