fyan1024/ScribFormer

[IEEE Transactions on Medical Imaging/TMI] This repository is the official implementation of the paper ScribFormer: Transformer Makes CNN Work Better for Scribble-based Medical Image Segmentation.

ScribFormer [TMI 2024]

This repository is the official implementation of the paper ScribFormer: Transformer Makes CNN Work Better for Scribble-based Medical Image Segmentation. Arxiv, ResearchGate, IEEEXplore

Datasets

ACDC

1. The ACDC dataset with mask annotations can be downloaded from ACDC.
2. The scribble annotations of ACDC have been released in ACDC scribbles.
3. The pre-processed ACDC data used for training could be directly downloaded from ACDC_dataset.

MSCMR

1. The MSCMR dataset with mask annotations can be downloaded from MSCMRseg.
2. The scribble annotations of MSCMRseg have been released in MSCMR_scribbles.
3. The scribble-annotated MSCMR dataset used for training could be directly downloaded from MSCMR_dataset.

Requirements

Some important required packages include:

• Python 3.8
• CUDA 11.3
• Pytorch 1.10.1.
• torchvision 0.11.2
• Some basic python packages such as Numpy, Scikit-image, SimpleITK, Scipy ......

Follow official guidance to install Pytorch.

Training

To train the model, run this command:

python train.py --root_path <dataset_path> --exp <path_to_save_model> --bilinear --linear_layer --fold MAAGfold --max_epoches 200 --weight_cam_subloss 0.25 0.5 0.75 1 --weight_cam 0.1

Evaluation

To evaluate the model, run this command:

python test.py --bilinear --linear_layer --fold MAAGfold --exp <path_to_save_model> --save_prediction

Citation

@article{li2023lvit,
  title={ScribFormer: Transformer Makes CNN Work Better for Scribble-based Medical Image Segmentation},
  author={Li, Zihan and Zheng, Yuan and Shan, Dandan and Yang, Shuzhou and Li, Qingde and Wang, Beizhan and Hong, Qingqi and Shen, Dinggang},
  journal={IEEE Transactions on Medical Imaging},
  year={2024},
  publisher={IEEE}
}