This repo holds code for TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

• [7/26/2024] TransUNet, which supports both 2D and 3D data and incorporates a Transformer encoder and decoder, has been featured in the journal Medical Image Analysis (link).

@article{chen2024transunet,
  title={TransUNet: Rethinking the U-Net architecture design for medical image segmentation through the lens of transformers},
  author={Chen, Jieneng and Mei, Jieru and Li, Xianhang and Lu, Yongyi and Yu, Qihang and Wei, Qingyue and Luo, Xiangde and Xie, Yutong and Adeli, Ehsan and Wang, Yan and others},
  journal={Medical Image Analysis},
  pages={103280},
  year={2024},
  publisher={Elsevier}
}

• [10/15/2023] 🔥 3D version of TransUNet is out! Our 3D TransUNet surpasses nn-UNet with 88.11% Dice score on the BTCV dataset and outperforms the top-1 solution in the BraTs 2021 challenge and secure the second place in BraTs 2023 challenge. Please take a look at the code and paper.

• Get models in this link: R50-ViT-B_16, ViT-B_16, ViT-L_16...

wget https://storage.googleapis.com/vit_models/imagenet21k/{MODEL_NAME}.npz &&
mkdir -p ../model/vit_checkpoint/imagenet21k &&
mv {MODEL_NAME}.npz ../model/vit_checkpoint/imagenet21k/{MODEL_NAME}.npz

Make sure the resulting file path is ../model/vit_checkpoint/imagenet21k/{MODEL_NAME}.npz relative to the repository root. Training will fail if this checkpoint is

missing. When using --vit_name R50-ViT-B_16, download the file named R50+ViT-B_16.npz.

All data are available so no need to send emails for data. Please use the BTCV preprocessed data and ACDC data.

Please prepare an environment with python=3.7, and then use the command "pip install -r requirements.txt" for the dependencies.

• Run the train script on synapse dataset. The batch size can be reduced to 12 or 6 to save memory (please also decrease the base_lr linearly), and both can reach similar performance.

CUDA_VISIBLE_DEVICES=0 python train.py --dataset Synapse --vit_name R50-ViT-B_16

• Run the test script on synapse dataset. It supports testing for both 2D images and 3D volumes.

python test.py --dataset Synapse --vit_name R50-ViT-B_16

• Train on the CBIS-DDSM dataset located at /data/xudosong/Transfer_Scratch/CBIS_max_last_withAug3_noise/CBIS_pre:

CUDA_VISIBLE_DEVICES=0,1 python train.py --dataset CBIS --vit_name R50-ViT-B_16 --img_size 256 --n_gpu 2

• Test on the CBIS-DDSM dataset:

python test.py --dataset CBIS --vit_name R50-ViT-B_16 --img_size 256

• Google ViT
• ViT-pytorch
• segmentation_models.pytorch

@article{chen2021transunet,
  title={TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation},
  author={Chen, Jieneng and Lu, Yongyi and Yu, Qihang and Luo, Xiangde and Adeli, Ehsan and Wang, Yan and Lu, Le and Yuille, Alan L., and Zhou, Yuyin},
  journal={arXiv preprint arXiv:2102.04306},
  year={2021}
}