obiyoag / BayeSeg

[MedIA Best Paper Award] Official implementation of MedIA paper "BayeSeg: Bayesian Modelling for Medical Image Segmentation with Interpretable Generalizability"

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

BayeSeg: Bayesian Modelling for Medical Image Segmentation with Interpretable Generalizability

WelcomeπŸ‘‹! This repository provides the official implementation of our paper BayeSeg: Bayesian Modelling for Medical Image Segmentation with Interpretable Generalizability [MedIA] [arXiv], which has been accepted by Medical Image Analysis.

BayeSeg framwork

πŸ” Details

Due to the cross-domain distribution shift aroused from diverse medical imaging systems, many deep learning segmentation methods fail to perform well on unseen data, which limits their real-world applicability. Recent works have shown the benefits of extracting domain-invariant representations on domain generalization. However, the interpretability of domain-invariant features remains a great challenge. To address this problem, we propose an interpretable Bayesian framework (BayeSeg) through Bayesian modeling of image and label statistics to enhance model generalizability for medical image segmentation. Specifically, we first decompose an image into a spatial-correlated variable and a spatial-variant variable, assigning hierarchical Bayesian priors to explicitly force them to model the domain-stable shape and domain-specific appearance information respectively. Then, we model the segmentation as a locally smooth variable only related to the shape. Finally, we develop a variational Bayesian framework to infer the posterior distributions of these explainable variables. The framework is implemented with neural networks, and thus is referred to as deep Bayesian segmentation. Quantitative and qualitative experimental results on prostate segmentation and cardiac segmentation tasks have shown the effectiveness of our proposed method. Moreover, we investigated the interpretability of BayeSeg by explaining the posteriors and analyzed certain factors that affect the generalization ability through further ablation studies.

πŸ“° News

πŸ“¦ Get started

Installation

  1. Create a virtual environment conda create -n BayeSeg python=3.10 and activate it conda activate BayeSeg
  2. Install PyTorch
  3. git clone https://github.com/obiyoag/BayeSeg
  4. Enter the BayeSeg folder cd BayeSeg and run pip install -r requirements.txt

Data preparation

  • The prostate dataset can be downloaded from here.
  • We trained on site RUNMC, which is splitted as
    RUNMC
    |-- test: [00, 08, 11, 16, 24, 28]
    |-- train: [01, 02, 03, 04, 06, 07, 09, 12, 13, 14, 15, 17, 18, 19, 20, 22, 23, 25, 26, 27, 29]
    `-- val: [05, 10, 21]
    

Training

python train.py --output_dir logs/BayeSeg --dataset_dir /path/to/dataset

Testing

python test.py --checkpoint_dir logs/BayeSeg --dataset_dir /path/to/dataset

Download Model

The trained model can be found here and should be placed at logs/BayeSeg/best_checkpoint.pth.

πŸ“’ To do

  • Prostate code release
  • Prostate model release
  • Cardiac code release
  • Cardiac model release
  • Tutorial

πŸ™‹β€β™€οΈ Feedback and Contact

πŸ›‘οΈ License

This project is under the Apache-2.0 license. See LICENSE for details.

πŸ™ Acknowledgement

  • We appreciate monai for developing the open-source deep learning framework for medical images.
  • We thank Quande Liu for summarizing the multi-site prostate dataset.

πŸ“ Citation

If our work is helpful in your research, please cite this as follows.

@article{Gao2023BayeSeg,
    author={Shangqi Gao, Hangqi Zhou, Yibo Gao, Xiahai Zhuang},
    title={BayeSeg: Bayesian Modelling for Medical Image Segmentation with Interpretable Generalizability},
    journal={Medical Image Analysis},
    year={2023}
}

@article{Gao2022BayeSeg,
    author={Shangqi Gao, Hangqi Zhou, Yibo Gao, Xiahai Zhuang},
    title={Joint Modeling of Image and Label Statistics for Enhancing Model Generalizability of Medical Image Segmentation},
    article={International Conference on Medical Image Computing and Computer Assisted Intervention},
    year={2022}
}

About

[MedIA Best Paper Award] Official implementation of MedIA paper "BayeSeg: Bayesian Modelling for Medical Image Segmentation with Interpretable Generalizability"

License:Apache License 2.0


Languages

Language:Python 100.0%