Official Pytorch Implementation for AttENT.
- 10/2021: Our paper is accepted as the regular paper of BIBM 2021 and will give a presentation online (15 mins). The relevant code is released here.
- 08/2021: Check out our paper [AttENT: Domain-Adaptive Medical Image Segmentation via Attention-Aware Translation and Adversarial Entropy Minimization] (submitted to BIBM 2021). With the aligned ensemble of attention-aware image pixel space and entropy-based feature space enables a well-trained segmentation modelto effectively transfer from source domain to target domain.
Publication | Paper | Slides
If you find this code useful for your research, please cite our work:
@inproceedings{li2021attent,
author={Li, Chen and Luo, Xin and Chen, Wei and He, Yulin and Wu, Mingfei and Tan, Yusong},
booktitle={2021 IEEE International Conference on Bioinformatics and Biomedicine (BIBM)},
title={AttENT: Domain-Adaptive Medical Image Segmentation via Attention-Aware Translation and Adversarial Entropy Minimization},
year={2021},
volume={},
number={},
pages={952-959},
doi={10.1109/BIBM52615.2021.9669620}}
}
- Hardware: PC with NVIDIA 1080T GPU. (others are alternative.)
- Software: Ubuntu 18.04, CUDA 10.0.130, pytorch 1.3.0, Python 3.6.9
- Package:
torchvision
tensorboardX
scikit-learn
glob
matplotlib
skimage
medpy
tqdm
- Clone the repo:
$ git clone https://github.com/lichen14/AttENT
$ cd AttENT
- Install OpenCV if you don't already have it:
$ conda install -c menpo opencv
- Install this repository and the dependencies using pip:
$ pip install -e <root_dir>
With this, you can edit the TriDL code on the fly and import function and classes of AttENT in other project as well. 3. Optional. To uninstall this package, run:
$ pip uninstall AttENT
- By default, the datasets are put in
<root_dir>/dataset
. - An alternative option is to explicitlly specify the parameters
DATA_DIRECTORY_SOURCE
andDATA_DIRECTORY_TARGET
in YML configuration files. - After download, The MMWHS dataset directory should have this basic structure:
- CHAOS: Please follow the instructions here to download images and semantic segmentation annotations. We extracted 20 labeled volumes from the T2-SPIR MRI training dataset.
- MALBCV: Please follow the instructions in here to download the images and ground-truths. We extracted 30 labeled volumes from the CT training dataset.
- The CHAOS and MALBCV datasets directory have the same structure but different content:
<root_dir>/data/CHAOS/mri_dataset/ % MRI samples root
<root_dir>/data/CHAOS/mri_dataset/train/image/ % MRI images
<root_dir>/data/CHAOS/mri_dataset/train/label/ % MRI annotation
<root_dir>/data/CHAOS/mri_list/ % MRI samples list
...
<root_dir>/data/MALBCV/ct_dataset/ % CT samples root
<root_dir>/data/MALBCV/ct_dataset/train/image/ % CT images
<root_dir>/data/MALBCV/ct_dataset/train/label/ % CT annotation
<root_dir>/data/MALBCV/ct_list/ % CT samples list
...
- For the common organs in the datasets, including liver, right kidney, left kidney, and spleen, we designed cross-modalities segmentation experiments between these organs.
- Initial pre-trained model can be downloaded from DeepLab-V2. Please download it and put it [AttENT/pretrained_models/].
The well-trained model can be downloaded here AttENT_deeplab. You can use the pre-trained model or your own model to make a test as following:
$ cd <root_dir>/advent
$ python test.py --cfg ./configs/<your_yml_name>.yml --exp-suffix <your_define_suffix>
- To ensure reproduction, the random seed has been fixed in the code. Still, you may need to train a few times to reach the comparable performance.
- The whole training process is not in the end-end manner, which means there is two stages (1.Attention-aware Image Alignment and 2.Adversarial Feature Alignment) with human intervention when training the AttENT. And, there are two separate training directions (MRI \rightarrow CT and CT \rightarrow MRI).
- By default, logs and snapshots are stored in
<root_dir>/experiments
with this structure:
<root_dir>/experiments/logs
<root_dir>/experiments/snapshots %output and trained model are stored in this file.
To perform cross-domain translation with Attention-aware CycleGAN:
$ python train.py --dataroot datasets/<dataset_name>/ --cuda
- We referred the implement of PyTorch-CycleGAN. Based on above work, we designed the Attention-aware CycleGAN, which is replaced by Generator5. following the instruction of PyTorch-CycleGAN, we revised the train.py with:
from models import Generator5
- With the Attention-aware CycleGAN, we obtained the translated multi-domain samples in the pixel space for the [2.Adversarial Feature Alignment in the Entropy Space].
To train AttENT separately from both directions:
$ cd AttENT/attent
$ python train_MRI2CT.py --cfg ./configs/<your_yml_name>.yml --exp-suffix <your_define_suffix> --tensorboard % using tensorboard
$ python train_CT2MRI.py --cfg ./configs/<your_yml_name>.yml --exp-suffix <your_define_suffix> --tensorboard % using tensorboard
This codebase is heavily borrowed from AdaptSegNet, PyTorch-CycleGAN and ADVENT. Thanks to following repos for sharing and we referred some of their codes to construct AttENT: