Enhancing Retinal Vessel Segmentation Generalization via Layout-Aware Generative Modelling

Installation

First, clone this repository and run install the environment:

cd RLAD
python -m venv RLAD_env
source RLAD_env/bin/activate
pip install --upgrade pip
pip install -e .

Data preparation

Download the following datasets from their official website: UZLF_TRAIN, GRAPE, MESSIDOR, PAPILA, MAGREB, ENRICH, 1000images,ddr_lesion_train, EYEPACS,G1020,idrid_lesion_train,ODIR_Train, UZLF_VAL and DRTiD.

Then organize them as follow:

├── RLAD
├── Databases
│   ├── Database_Name_1
│   │   ├── images
│   │   │   ├── image_name1.png   
│   │   │   ├── image_name2.png
│   │   │   ... 
│   ├── Database_Name_2
│   │   ├── images
│   │   │   ├── image_name1.png   
│   │   │   ├── image_name2.png
│   │   │   ...
│   ├── Database_Name_3
│   │   ├── images
...

Training

Run the following command: (Original RLAD model have been trained using 4 A100-40gb GPUs).

cd RLAD/Runs
bash train.sh

Generate Images

To customize the image generation process, modify the following fields in configs/configsDiT/RLAD.yaml:

• load_weights_from: Set the path of the checkpoint you want to use (default is the model used in RLAD paper)
• CD_cond: Set to True if you want the model to use optic cup and disc conditioning during generation. Set to False to allow the model to decide freely, without this conditioning.
• L_cond: Set to True to provide lesion conditioning for the generation process. Set to False to let the model generate images without lesion conditioning.
• n_gen_per_samples: The number of different sample of the same conditioning to generate

After updating the configuration file, generate images by running:

cd RLAD/Runs
bash generate_images.sh

Compute FID and RET-FD

After generating some augmented images for a given test dataset (see the configs/configsDiT/RLAD.yaml) run:

cd RLAD
source RLAD_env/bin/activate
python Scripts/compute_FD.py --yaml configs/configsDiT/RLAD.yaml --bs 160

Main Results

Evaluation of Realism

Gen Model	Conditioning	FID ↓	RET-FD ↓
StyleGAN [1]	L	138.0	120.8
StyleGAN2 [2]	Demographics	98.1	116.0
StyleGAN2† [3]	AV	122.8	-
Pix2PixHD† [3]	AV	86.8	-
RLAD (Ours)	AV + L + CD	30.3	79.7

Table: Realism of Generated Images. Lower FID and RET-FD on the DRTiD dataset indicate closer alignment with real data, reflecting realism. Notably, RLAD is able to generate controllable and more realistic retinal images. Models marked ^† were trained and evaluated on private data.

Improved performance using RLAD as random data augmentation

(with no CD and Lesion conditioning during generation and 15 generated image per sample)

Backbone	Avg In-Domain	Avg Near-Domain	Avg Out-of-Domain
Little W-Net [4]	-	67.9	45.5
Automorph [5]	80.2	71.7†	57.9†
VascX [6]	81.2	76.0	60.5
LUNet [7]	83.4	77.3	61.1
DinoV2small [8]	82.0	76.6	64.0
+ RLAD (Ours)	82.3	77.5	66.6
RETFound [9]	81.8	76.9	65.2
+ RLAD (Ours)	83.4	79.9	69.9
SwinV2tiny [10]	83.1	79.6	68.9
+ RLAD (Ours)	83.3	79.9	70.8
SwinV2large [10]	83.4	80.5	72.1
+ RLAD (Ours)	83.4	80.7	72.5

Table: Average Dice scores for artery and vein segmentation across In-Domain, Near-Domain, and Out-of-Domain datasets. RLAD denotes results with our method. ^† indicates data leakage during training.

References

1. Benjamin Hou, Amir Alansary, Daniel Rueckert, and Bernhard Kainz. High-fidelity diabetic retina fundus image synthesis from freestyle lesion maps, 2022.
2. Sarah Müller, Lisa M. Koch, P. A. Lensch, Hendrik, and Philipp Berens. Disentangling representations of retinal images with generative models, 2024.
3. Sojung Go, Younghoon Ji, Sang Jun Park, and Soochahn Lee. Generation of structurally realistic retinal fundus images with diffusion models. In IEEE/CVF Conference on Computer Vision and Pattern Recogni- tion Workshops (CVPRW), pages 2335–2344, Seattle, WA, USA, 2024.
4. Adrian Galdran, André Anjos, José Dolz, Hadi Chakor, Hervé Lombaert, and Ismail Ben Ayed. State-of-the- art retinal vessel segmentation with minimalistic models. Scientific Reports, 12(1):6174, 2022.
5. Yukun Zhou, Siegfried K Wagner, Mark A Chia, An Zhao, Moucheng Xu, Robbert Struyven, Daniel C Alexander, Pearse A Keane, and others. AutoMorph: automated retinal vascular morphology quantification via a deep learning pipeline. Translational Vision Science & Technology, 11(7):12, 2022.
6. Jose Vargas Quiros, Bart Liefers, Karin van Garderen, Jeroen Vermeulen, Eyened Reading Center, Sinergia Consortium, and Caroline Klaver. VascX models: Model ensembles for retinal vascular analysis from color fundus images, 2024.
7. Jonathan Fhima, Jan Van Eijgen, Marie-Isaline Billen Moulin-Romsée, Heloı¨se Brackenier, Hana Ku- lenovic, Valérie Debeuf, Marie Vangilbergen, Moti Freiman, Ingeborg Stalmans, and Joachim A Be- har. LUNet: deep learning for the segmentation of arterioles and venules in high resolution fundus images. Physiological Measurement, 45(5):055002, 2024.
8. Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, and others. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023.
9. Yukun Zhou, Mark A Chia, Siegfried K Wagner, Murat S Ayhan, Dominic J Williamson, Robbert R Struyven, Timing Liu, Moucheng Xu, Mateo G Lozano, Peter Woodward-Court, and others. A foundation model for generalizable disease detection from retinal images. Nature, 622(7981):156– 163, 2023.
10. Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, and others. Swin transformer v2: Scaling up capacity and resolution. In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12009–12019, 2022.

LICENSE

This project is licensed under the Creative Commons Attribution-NonCommercial 4.0 International License, see LICENSE file, which prohibits commercial use.