/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 (MedIA2022) and WSL4MIS (MICCAI2022). 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={Medical Image Computing and Computer Assisted Intervention -- MICCAI 2022},
  year={2022},
  pages={528--538}}
  
  @article{luo2022word,
  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, Jieneng Chen, Tao Song, Xiaofan Zhang, Kang Li, Dimitris N. Metaxas, Guotai Wang, and Shaoting Zhang},
  journal={Medical Image Analysis},
  volume={82},
  pages={102642},
  year={2022},
  publisher={Elsevier}}
  
  @misc{wsl4mis2020,
  title={{WSL4MIS}},
  author={Luo, Xiangde},
  howpublished={\url{https://github.com/Luoxd1996/WSL4MIS}},
  year={2021}}

  • A re-implementation of this work based on the PyMIC can be found here (WSLDMPLS).

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.
  • The ISBI-MR-Prostate-2013 dataset with mask annotation can be downloaded from TCIA, the scribble annotation of this dataset (annotated by Mr. X. Luo and M.D. W. Liao) can be downloaded GoogleDrive and BaiduPan.
  • 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/HiLab-git/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.