TransDeepLab

The official code for "TransDeepLab: Convolution-Free Transformer-based DeepLab v3+ for Medical Image Segmentation".

Updates

• July 19, 2022: Initial release.

• July 4, 2022: Submitted to MICCAI PRIME2022 [Under Review] [Accepted].

Setting up and Training

• We use the code base from the Swin-Unet GitHub repo as our starting point.

• In order to run the code and experiments, you need to first install the dependencies and then download and move the data to the right place.

  • For the Synapse dataset, we used the data provided by TransUnet's authors.
  • For ISIC 2017-18 datasets, we used the ISIC Challenge datasets link.
  • For the PH² dataset, we used this link.

• We have put the required instructions for doing the above steps in the ./setup.sh file in the repo for your convenience. cd to this repo directory and then run it to install dependencies and download and move data to the right dir.

• Download Swin-T pre-trained weights from this link and put it into the folder pretrained_ckpt/.

• Then you need to run the train.py file with the appropriate arguments to run your experiment. let's just see an example for now for the Synapse dataset:

python train.py --config_file 'swin_224_7_{# of SSPP}level' --dataset Synapse --root_path './data' --max_epochs 200 --output_dir '.'  --img_size 224 --base_lr 0.05 --batch_size 24

Config files and hyperparameters

• All the hyperparameters related to building different models are in separate files in ./model/configs directory.

• For each experiment, you need to make your desired config_name.py file and put it in the model/configs dir and then enter the file name (without .py suffix) in the python train.py command you saw in the previous section after --config_file arg.

• For the rest of the hyperparameters like batch_size, max_epochs, base_lr, ... look at the Swin-Unet or the code here to see what you can change and how to do so.

Test

• The model can be tested with the following command using test.py file. During training, model checkpoints will be saved to disk with the following format: output_dir/{config_file_name}_epoch_{epoch_num}.pth.

• It takes the checkpoint (model weight file) name as an input argument and loads the appropriate config file from the configs dir.

• Other arguments and flags can be given to the test.py file if some settings need to be modified but --ckpt_path and --config_file are the only required arguments.

• Trained weights for our best-reported results in the paper for the Synapse dataset are easily accessible from this link, where you could download it as a sole folder via gdown or setting specific links listed under the below table:

• :exclamation: Remember to put these weights in a specific folder that you are going to address them with test.py via --ckpt_path flag.

• Comparison results table on the Synapse dataset:

• Impact of modifying modules inside the proposed method:

• A look at the number of parameters:

Visualization

• Results on the Synapse dataset:

Image	Ground Truth	Prediction

References

• TransUNet

• Swin-Unet

  Citation

  @article{azad2022transdeeplab,
  title={TransDeepLab: Convolution-Free Transformer-based DeepLab v3+ for Medical Image Segmentation},
  author={Azad, Reza and Heidari, Moein and Shariatnia, Moein and Aghdam, Ehsan Khodapanah and Karimijafarbigloo, Sanaz and Adeli, Ehsan and Merhof, Dorit},
  journal={arXiv preprint arXiv:2208.00713},
  year={2022}
  }