My modified version that fixes the issues of the original repo to run it now with python3.10 at 2023
3D U2-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.
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 segmentation task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a single 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, regardless 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 sharing cross-channel correlations which can be modeled with pointwise convolution. We evaluate the 3D U2-Net on five organ segmentation datasets. Experimental 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.
Brief instructions to apply the code:
-
My implementation uses
python3.10.12
-
Clone the repository by:
!git clone https://github.com/arrafi-musabbir/unet_3D_multi_domain_medical_image.git
-
Create a virtual env with
venv
(if you're running from local machine)!python -m venv venv
-
first, install all the necessary requirements
!pip install -r requirements.txt
-
Download the datasets from Medical Segmentation Decathlon
-
Please put the downloaded files in the
dataset
directory!tar -xf /content/drive/MyDrive/MSD/Task02_Heart.tar -C /content/3D-Universal-U-Net-for-Multi-Domain-Medical-Image-Segmentation/dataset
-
scripts in
u2net_torch_src
are utilized as below:data_explore.py
is to explore the characteristics of the images, e.g. pixel spacings.
!python data_explore.py
preprocess_taskSep.py
is used to do offline preprocessing (e.g. cropping, resampling) of the data samples to save time for training.
usage: preprocess_taskSep.py [-h] [--nProc NPROC] [--tasks TASKS [TASKS ...]]
!python preprocess_taskSep.py --tasks Task02_Heart Task04_Hippocampus
train_model_no_adapters.py
is the main file to train the independent models as well as the shared model.
usage: train_model_no_adapters.py [-h] [--tasks TASKS [TASKS ...]] [--trainMode TRAINMODE] [--module MODULE] [--ckp CKP] [--resume_ckp RESUME_CKP] [--resume_epoch RESUME_EPOCH] [--fold FOLD] [--model MODEL] [--out_tag OUT_TAG] [--base_outChans BASE_OUTCHANS] [--predict]
!python train_model_no_adapters.py --module series_adapter --tasks Task02_Heart Task04_Hippocampus
train_model_wt_adapters.py
is the main file to train the proposed universal model with separable convolution.
usage: train_model_wt_adapters.py [-h] [--tasks TASKS [TASKS ...]] [--trainMode TRAINMODE] [--module MODULE] [--ckp CKP] [--resume_ckp RESUME_CKP] [--resume_epoch RESUME_EPOCH] [--fold FOLD] [--model MODEL] [--out_tag OUT_TAG] [--base_outChans BASE_OUTCHANS] [--predict PREDICT] [--pred_epoch PRED_EPOCH]
!python train_model_wt_adapters.py --module parallel_adapter --tasks Task02_Heart Task04_Hippocampus
-
Terminal commands to train all models are presented in
train_models.sh
.
To accelerate training, the authors built a fast tool to do online image augmentation with CUDA on GPU(especially for elastic deformation). cuda_spatial_defrom.
If you use this code, please cite their 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
The authors 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.