Joint Modeling of Image and Label Statistics for Enhancing Model Generalizability of Medical Image Segmentation

TL;DR

This paper introduces a deep Bayesian framework that models image and label statistics jointly, improving the generalizability of medical image segmentation across different data domains.

Contribution

It proposes a novel deep Bayesian segmentation method that leverages domain-irrelevant contours to enhance model robustness and generalization in medical imaging.

Findings

Achieved state-of-the-art generalization in cross-sequence cardiac MRI segmentation.

Outperformed existing models with over 0.47 higher average Dice score.

Demonstrated strong transferability from LGE MRI to T2 images.

Abstract

Although supervised deep-learning has achieved promising performance in medical image segmentation, many methods cannot generalize well on unseen data, limiting their real-world applicability. To address this problem, we propose a deep learning-based Bayesian framework, which jointly models image and label statistics, utilizing the domain-irrelevant contour of a medical image for segmentation. Specifically, we first decompose an image into components of contour and basis. Then, we model the expected label as a variable only related to the contour. Finally, we develop a variational Bayesian framework to infer the posterior distributions of these variables, including the contour, the basis, and the label. The framework is implemented with neural networks, thus is referred to as deep Bayesian segmentation. Results on the task of cross-sequence cardiac MRI segmentation show that our method set a new state of the art for model generalizability. Particularly, the BayeSeg model trained with LGE MRI generalized well on T2 images and outperformed other models with great margins, i.e., over 0.47 in terms of average Dice. Our code is available at https://zmiclab.github.io/projects.html.