Anatomically-aware Uncertainty for Semi-supervised Image Segmentation

MedIA 2024 [paper], MICCAI 2022 [paper] [presentation] [poster]

TL;DR: A novel uncertainty estimation using anatomically-aware representation to guide the segmentation model in a low-data regime.

Keywords: Semi-Supervised learning; Anatomically-aware Representation; Labeling Representation; Image Segmentation; Uncertainty.

Dependencies

This code depends on the following libraries:

• Pytorch (1.8.0+cu111)
• Python >= 3.8
• tensorboardX
• some basic libraries: numpy, glob, skimage, matplotlib, tqdm...

Datasets

• LA, 2018
• FLARE, 2021

Training

Training of our approach involves two steps:

1. DAE (Denoising Autoencoder) model training with available labels

cd code_DAE
python train_DAE.py --exp DAE_L10 --nb_labels 26 --total_labels 260 --emb_dim 512

2. Segmentation model training with DAE under limited labels

cd code_DAE
python train_Abdomen_meanteacher_DAE_certainty.py --exp L10_r1 --nb_labels 26 --total_labels 260 --model_DAE 'DAE_L10/model.pth' --emb_dim 512

Testing

Testing segmentation results

python test_Abdomen.py --model 'L10_r1/best_model.pth'

Citation

Please cite our paper if you find this code or our work useful for your research.

@article{adiga2023anatomically,
  title={Anatomically-aware Uncertainty for Semi-supervised Image Segmentation},
  author={Adiga V, Sukesh and Dolz, Jose and Lombaert, Herve},
  journal={Medical Image Analysis (MedIA)},
  year={2024}
}

@article{adiga2022leveraging,
  title={Leveraging Labeling Representations in Uncertainty-based Semi-supervised Segmentation},
  author={Adiga V, Sukesh and Dolz, Jose and Lombaert, Herve},
  journal={Medical Image Computing and Computer Assisted Intervention (MICCAI)},
  year={2022}
}

References

• Uncertainty-aware Self-ensembling Model for Semi-supervised (UAMT) [paper][code]
• Semi-supervised Learning for Medical Image Segmentation (SSL4MIS) [paper][code]

Any questions?

For more information, please get in touch with Sukesh Adiga ([email protected]).

License

This project is licensed under the terms of the MIT license.