AADG

By Junyan Lyu, Yiqi Zhang, Yijin Huang, Li Lin, Pujin Cheng, Xiaoying Tang.

This repository contains an official implementation of AADG for the TMI paper "AADG: Automatic Augmentation for Domain Generalization on Retinal Image Segmentation".

Quick start

Environment

This code is developed using on Python 3.8.5 and Pytorch 1.8.0 on CentOS 7 with NVIDIA GPUs. Training and testing are performed using 1 Tesla A100 GPU with CUDA 11.1. Other platforms or GPUs are not fully tested.

Install

Install Pytorch
Install dependencies

pip install -r requirements.txt

Replace <ENV>/lib/python3.8/site-packages/segmentation_models_pytorch/base/heads.py in your python environment with models/heads.py provided in this repository.
Make sure your gcc, cmake and cuda versions are compatitable with pykeops.

Data

Make a dataset directory.

cd AADG
mkdir dataset

Download the OD/OC datasets into dataset.
Download the retinal vessel datasets into dataset.
Your dataset directory should look like this:

AADG
-- dataset
   |-- RVS
   |   |-- CHASEDB1
   |   |-- DRIVE
   |   |-- HRF
   |   |-- STARE
   |-- Fundus
   |   |-- Domain1
   |   |-- Domain2
   |   |-- Domain3
   |   |-- Domain4

Train

Please specify the configuration file in experiments.

python run.py --cfg <CONFIG-FILE> --output_dir <CUSTOM-OUTPUT-DIR>

For example,

python run.py --cfg experiments/rvs_sinkhorn/diversity_ex.yaml --output_dir output/

Citation

If you find this repository useful, please consider citing AADG paper:

@ARTICLE{9837077,
  author={Lyu, Junyan and Zhang, Yiqi and Huang, Yijin and Lin, Li and Cheng, Pujin and Tang, Xiaoying},
  journal={IEEE Transactions on Medical Imaging}, 
  title={AADG: Automatic Augmentation for Domain Generalization on Retinal Image Segmentation}, 
  year={2022},
  volume={},
  number={},
  pages={1-1},
  doi={10.1109/TMI.2022.3193146}}