Kalrfou / GCtx-UNet

MIT License
20 stars 1 forks source link

GCtx-UNet: Efficient Deep Network for Medical Image Segmentation

GCtx-UNet is a U-shaped network architecture that incorporates the Global Context Vision Transformer (GC-ViT) to enhance medical image segmentation by effectively capturing both global and local features. Feel free to check out our preprint on arXiv. ⁣

Pre-trained model:

1. Download pre-trained GC-ViT transformer model (GCViT-xxtiny) pre-trained on ImageNet1K:

[Get pre-trained model in this link] (https://drive.usercontent.google.com/download?id=1apSIWQCa5VhWLJws8ugMTuyKzyayw4Eh&export=download&authuser=0): Put pretrained xx-Tiny into folder "pretrained_ckpt/"

2. Download pre-trained GC-ViT transformer model (GCViT-xxtiny) pre-trained on MedNet:

Train/Test

1) Train

CUDA_VISIBLE_DEVICES=0  python -W ignore train.py 

2) Testing

Get pre-trained GCtx-UNet model weights for the Synapse dataset: link

download the file and put it into the folder model_out.

CUDA_VISIBLE_DEVICES=0  python -W ignore test.py --is_saveni

References

Citation


@article{alrfou2024gctx,
  title={GCtx-UNet: Efficient Network for Medical Image Segmentation},
  author={Alrfou, Khaled and Zhao, Tian},
  journal={arXiv preprint arXiv:2406.05891},
  year={2024}
}