3D U²-Net: A 3D Universal U-Net for Multi-Domain Medical Image Segmentation

by Chao Huang, Qingsong Yao, Hu Han, Shankuan Zhu, Shaohua Zhou. This is a code repo of the paper early accepted by MICCAI2019.

In case of any questions about this repo, please feel free to contact Chao Huang(huangchao09@zju.edu.cn).

Abstract. Fully convolutional neural networks like U-Net have been the state-of-art methods in medical image segmentation. Practically, a network is highly specialized and trained separately for each segmenta- tion task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a sin- gle model with the addition of a minimal number of parameters to steer to each task. Inspired by the recent success of multi-domain learning in image classification, for the first time we explore a promising universal architecture that can handle multiple medical segmentation tasks, re- gardless of different organs and imaging modalities. Our 3D Universal U-Net (3D U2-Net) is built upon separable convolution, assuming that images from different domains have domain-specific spatial correlations which can be probed with channel-wise convolution while also share cross- channel correlations which can be modeled with pointwise convolution. We evaluate the 3D U2-Net on five organ segmentation datasets. Experimen- tal results show that this universal network is capable of competing with traditional models in terms of segmentation accuracy, while requiring only 1% of the parameters. Additionally, we observe that the architecture can be easily and effectively adapted to a new domain without sacrificing performance in the domains used to learn the shared parameterization of the universal network.

Overview

Brief instruction to apply the code:

1. Requirements are listed in requirments.txt.
2. Please put the datasets downloaded from Medical Segmentation Decathlon in dataset.
3. data_explore.py is to explore the characteristics of the images, e.g. pixel spacings.
4. preprocess_taskSep.py is used to do offline preprocessing (e.g. cropping, resampling) of the data samples to save time for training.
5. train_model_no_adapters.py is the mainfile to train the independent models as well as the shared model.
6. train_model_wt_adapters.py is the mainfile to train the propsed universal model with separable convolution.
7. Terminal commands to train all models are presented in train_models.sh.

To accelerate training, we built a fast tool to do online image augmentation with CUDA on GPU(especially for elastic deformation). cuda_spatial_defrom.

Citation

If you use this code, please cite our paper as:

Huang C., Han H., Yao Q., Zhu S., Zhou S.K. (2019) 3D U 2-Net: A 3D Universal U-Net for Multi-domain Medical Image Segmentation. In: Shen D. et al. (eds) Medical Image Computing and Computer Assisted Intervention – MICCAI 2019. MICCAI 2019. Lecture Notes in Computer Science, vol 11765. Springer, Cham

Acknowledgement

We give a lot of thanks to the open-access data science community for the public data science knowledge. Special thanks are given to @Guotai Wang and @Rebuffi as some of the code is borrowed from their repos: https://github.com/taigw/brats17/ and https://github.com/srebuffi/residual_adapters.