A pre-trained model for volumetric (3D) multi-organ segmentation from CT image.

A pre-trained Swin UNETR [1,2] for volumetric (3D) multi-organ segmentation using CT images from Beyond the Cranial Vault (BTCV) Segmentation Challenge dataset [3].

The training data is from the BTCV dataset (Please regist in Synapse and download the Abdomen/RawData.zip). The dataset format needs to be redefined using the following commands:

unzip RawData.zip
mv RawData/Training/img/ RawData/imagesTr
mv RawData/Training/label/ RawData/labelsTr
mv RawData/Testing/img/ RawData/imagesTs

The training was performed with at least 32GB-memory GPUs.

Actual Model Input: 96 x 96 x 96

Input: 1 channel CT image

Output: 14 channels: 0:Background, 1:Spleen, 2:Right Kidney, 3:Left Kideny, 4:Gallbladder, 5:Esophagus, 6:Liver, 7:Stomach, 8:Aorta, 9:IVC, 10:Portal and Splenic Veins, 11:Pancreas, 12:Right adrenal gland, 13:Left adrenal gland

A graph showing the validation mean Dice for 5000 epochs.

This model achieves the following Dice score on the validation data (our own split from the training dataset):

Mean Dice = 0.8283

Note that mean dice is computed in the original spacing of the input data.

Execute training:

Override the train config to execute multi-GPU training:

Override the train config to execute evaluation with the trained model:

Execute inference:

Export checkpoint to TorchScript file:

TorchScript conversion is currently not supported.

This is an example, not to be used for diagnostic purposes.

[1] Hatamizadeh, Ali, et al. "Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images." arXiv preprint arXiv:2201.01266 (2022). https://arxiv.org/abs/2201.01266.

[2] Tang, Yucheng, et al. "Self-supervised pre-training of swin transformers for 3d medical image analysis." arXiv preprint arXiv:2111.14791 (2021). https://arxiv.org/abs/2111.14791.

[3] Landman B, et al. "MICCAI multi-atlas labeling beyond the cranial vault–workshop and challenge." In Proc. of the MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge 2015 Oct (Vol. 5, p. 12).