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

This project was originally developed for our previous works WORD[Paper & Dataset]. We also investigated the comprehensive results on the ACDC dataset on this Tech Report. If you use this codebase 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={arXiv preprint arXiv:2203.02106},
  year={2022}}

  @article{luo2021word,
  title={{WORD}: Revisiting Organs Segmentation in the Whole Abdominal Region},
  author={Luo, Xiangde and Liao, Wenjun and Xiao, Jianghong and Song, Tao and Zhang, Xiaofan and Li, Kang and Wang, Guotai and Zhang, Shaoting},
  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.

felixshiyong/WSL4MIS