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

Clone this project.

git clone https://github.com/Luoxd1996/WSL4MIS
cd WSL4MIS

Data pre-processing os used or the processed data.

cd code
python dataloaders/acdc_data_processing.py

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).

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

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.

DunkDong/WSL4MIS