TranSegNet

This repo holds code for [https://pubmed.ncbi.nlm.nih.gov/37109505/]

Inspired by TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation

📰 News

• [10/15/2023] 🔥 3D version of TransUNet is out! 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. Please take a look at the code and paper.

Usage

1. Download Google pre-trained ViT models

• 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 ../model/vit_checkpoint/imagenet21k &&
mv {MODEL_NAME}.npz ../model/vit_checkpoint/imagenet21k/{MODEL_NAME}.npz

2. Prepare data

Please go to "./datasets/README.md" for details, or use the preprocessed data and data2 for research purposes.

3. Environment

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

4. Train/Test

• 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

5. FYI

This repo is NOT a official code base of papers above. This is a repo that contains reversely built model based on the thesis of TranSegnet. Made for the cataract detection experiment in the fundus photograph.

Made by Maison.Jang, Samsung Medical Center AI Reasearch

Reference

• Google ViT
• ViT-pytorch
• segmentation_models.pytorch

Citations

@article{jhang2023transegnet,
  title={TranSegNet: Hybrid CNN-Vision Transformers Encoder for Retina Segmentation of Optical Coherence Tomography},
  author={Zhang, Y.; Li, Z.; Nan, N.; Wang, X},
  journal={MDPI Life 2023}
  year={2023}
@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}
}