/WSL4MIS

Scribbles or Points-based weakly-supervised learning for medical image segmentation, a strong baseline, and tutorial for research and application.

Primary LanguagePythonMIT LicenseMIT

Weakly-supervised learning for medical image segmentation (WSL4MIS).

  • This project was originally developed for our two previous works WORD and WSL4MIS (MICCAI2022, early accept & student travel award). If you use this project in your research, please cite the following works:

      @article{luo2022scribbleseg,
      title={Scribble-Supervised Medical Image Segmentation via Dual-Branch Network and Dynamically Mixed Pseudo Labels Supervision},
      author={Xiangde Luo, Minhao Hu, Wenjun Liao, Shuwei Zhai, Tao Song, Guotai Wang, Shaoting Zhang},
      journal={MICCAI},
      year={2022}}
      
      @article{luo2021word,
      title={{WORD}: A large scale dataset, benchmark and clinical applicable study for abdominal organ segmentation from CT image},
      author={Xiangde Luo, Wenjun Liao, Jianghong Xiao, Tao Song, Xiaofan Zhang, Kang Li, Dimitris N. Metaxas, Guotai Wang, and Shaoting Zhang},
      journal={arXiv preprint arXiv:2111.02403},
      year={2021}}
      
      @misc{wsl4mis2020,
      title={{WSL4MIS}},
      author={Luo, Xiangde},
      howpublished={\url{https://github.com/Luoxd1996/WSL4MIS}},
      year={2021}}
    

Dataset

  • The ACDC dataset with mask annotations can be downloaded from: ACDC.
  • The Scribble annotations of ACDC can be downloaded from: Scribble.
  • The data processing code in Here the pre-processed ACDC data in Here.
  • To simulate the scribble annotation for other datasets, we further provide the simulation code at Here.

Requirements

Some important required packages include:

  • Pytorch version >=0.4.1.
  • TensorBoardX
  • Python >= 3.6
  • Efficientnet-Pytorch pip install efficientnet_pytorch
  • Some basic python packages such as Numpy, Scikit-image, SimpleITK, Scipy ......

Follow official guidance to install Pytorch.

Usage

  1. Clone this project.
git clone https://github.com/Luoxd1996/WSL4MIS
cd WSL4MIS
  1. Data pre-processing os used or the processed data.
cd code
python dataloaders/acdc_data_processing.py
  1. Train the model
cd code
bash train_wss.sh # train model with scribble or dense annotations.
bash train_ssl.sh  # train model with mix-supervision (mask annotations and without annotation).
  1. Test the model
python test_2D_fully.py --sup_type scribble/label --exp ACDC/the trained model fold --model unet
python test_2D_fully_sps.py --sup_type scribble --exp ACDC/the trained model fold --model unet_cct
  1. Training curves on the fold1: Note: pCE means partially cross-entropy, TV means total variation, label denotes supervised by mask, scribble represents just supervised by scribbles.

Implemented methods

Acknowledgement

  • The GatedCRFLoss is adapted from GatedCRFLoss for medical image segmentation.
  • The codebase is adapted from our previous work SSL4MIS.
  • The WORD dataset will be presented at WORD.