Skip to content

Dynamically Mixed Soft Pseudo-label Supervision for Scribble-Supervised Medical Image Segmentation

Notifications You must be signed in to change notification settings

HiLab-git/DMSPS

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

30 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

DMSPS

official code for: DMSPS: Dynamically mixed soft pseudo-label supervision for scribble-supervised medical image segmentation. MedIA 2024 MedIA. And the previous version is published on the MICCAI 2022.

Overall Framework

The overall framework of DMSPS: overall

Citation

If you use this project in your research, please cite the following works:

@article{han2024dmsps,
  title={DMSPS: Dynamically mixed soft pseudo-label supervision for scribble-supervised medical image segmentation},
  author={Han, Meng and Luo, Xiangde and Xie, Xiangjiang and Liao, Wenjun and Zhang, Shichuan and Song, Tao and Wang, Guotai and Zhang, Shaoting},
  journal={Medical Image Analysis},
  pages={103274},
  year={2024},
  publisher={Elsevier}
}

@inproceedings{luo2022scribble,
  title={Scribble-supervised medical image segmentation via dual-branch network and dynamically mixed pseudo labels supervision},
  author={Luo, Xiangde and Hu, Minhao and Liao, Wenjun and Zhai, Shuwei and Song, Tao and Wang, Guotai and Zhang, Shaoting},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={528--538},
  year={2022},
  organization={Springer}
}

Dataset

  • The ACDC dataset with mask annotations can be downloaded from: ACDC.
  • The Scribble annotations of ACDC can be downloaded from: Scribble.
  • The ACDC data can also be downloaded from the following BaiduPan link, where the same training set, test set and verification set as in DMSPS have been divided: ACDC BaiduPan. The extraction code is:et38 .
  • The WORD dataset can be downloaded from WORD.
  • The BraTS2020 dataset can be downloaded from BraTS2020. Note that this work aimed to segment two foreground classes: the tumor core and the peritumoral edema. Scribbles could be genearted by this simulation code: scribble-generate

Usage: [Taking the ACDC segmentation task as an example]

Step0:

  1. Clone this project.
git clone https://github.com/HiLab-git/DMSPS
cd DMSPS
  1. Data pre-processing.
cd code/dataloaders
python preprocess_ACDC.py

Quick test using pre-trained checkpoints:

  1. the first stage
cd code/test
python test_2d_forall_fast_txtver.py --data_name Heart_ACDC_Example \
--exp A_weakly_SPS_2d --fold stage1 --model unet_cct --tt_num 1
  1. the second stage
python test_2d_forall_fast_txtver.py --data_name Heart_ACDC_Example \
--exp A_weakly_SPS_2d --fold stage2 --model unet_cct --tt_num 1

Train and test for the first stage

  1. Train the model
cd code/train
python A_train_weaklySup_SPS_2d_soft.py
  1. Test the model
cd code/test
python test_2d_forall_fast_txtver.py 

Train and test for the second stage

  1. test on trainSet and get the uncertainty-filterd pseudo-label
cd code/test
python test_2d_forall_fast_txtver_forTrainSetUncertaintyOnly_Mean.py \
    --data_root_path $yourPath/ACDC2017/ACDC_for2D \
    --model unet_cct --exp A_weakly_SPS_2d --fold stage1 --threshold 0.1 --tt_num 3
  1. deal with the produced confident expanded annotation into h5 file and get the txt
cd code/dataloader
python retrain_postProcess_ACDC_uncertainty.py \
    --data_root_path  $yourPath/ACDC2017/ACDC \
    --func 0 --txtName trainReT01
python retrain_postProcess_ACDC_uncertainty.py \
    --data_root_path  $yourPath/ACDC2017/ACDC \
    --func 1 --txtName trainReT01
  1. train for stage2
cd code/train
python A_train_weaklySup_SPS_2d_soft_retrainUncertainty.py \
    --data_root_path  $yourPath/ACDC2017/ACDC_for2D \
    --model unet_cct --exp A_weakly_SPS_2d --fold stage2 \
    --sup_type pseudoLab --trainData trainReT01.txt
  1. test for stage2
cd code/test
python test_2d_forall_fast_txtver.py \
    --data_root_path $yourPath/ACDC2017/ACDC_for2D \
    --model unet_cct --exp A_weakly_SPS_2d --fold stage2

Acknowledgement

The code of scribble-supervised learning framework is borrowed from WSL4MIS

About

Dynamically Mixed Soft Pseudo-label Supervision for Scribble-Supervised Medical Image Segmentation

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published