Diffusion semantic segmentation model: A generative model for medical image segmentation based on joint distribution

Background

The mainstream semantic segmentation schemes in medical image segmentation are essentially discriminative paradigms based on conditional distributions $p\left(class|feature\right)$. Although efficient and straightforward, this prevalent paradigm focuses solely on extracting image features while ignoring the underlying data distribution $p\left(feature|class\right)$. Therefore, the learned feature space exhibits inherent instability, which directly affects the precision of the model in delineating anatomical boundaries.

Purpose

This paper reformulates the semantic segmentation task as a distribution alignment problem for medical image segmentation, aiming to minimize the gap between model predictions and ground truth labels by modeling the joint distribution of the data.

Methods

We propose a novel segmentation architecture based on joint distribution, called Denoising Semantic Segmentation Model (DSSM). We propose learning classification decision boundaries in pixel feature space and modeling joint distributions in latent feature space. Specifically, DSSM optimizes probability maps based on pixel feature classification through Bayesian posterior probabilities. To this end, we design a Feature Fusion Module (FFM) to guide the generative module in inference and provide label features for the semantic module. Furthermore, we introduce a stable Markov inference process to reduce inference offset. Finally, the joint distribution-based model is end-to-end trained in a discriminative manner, that is, maximizing $p\left(class|feature\right)$, which endows DSSM with the strengths of both generative and discriminative models.

Results

The image datasets utilized in this study are from different modalities, including MRI scans, x-ray images, and skin lesion photographic images, demonstrating superior performance compared to state-of-the-art (SOTA) discriminative models. Specifically, DSSM achieved a Dice coefficient of 0.8871 in MSD cardiac MRI segmentation, 0.9451 in ACDC left ventricular MRI segmentation, and 0.9647 in x-ray image segmentation. DSSM also reached 0.8731 Dice in prostate MRI segmentation. Furthermore, in the field of skin lesion segmentation, DSSM achieved a Dice score of 0.8869 on the ISIC 2018 dataset and delivered exceptional performance with 0.9421 on the PH2 dataset. Besides the Dice score, HD95, mIoU, Precision, and Recall are evaluated across the above datasets, which further demonstrate the superior performance of DSSM.

Conclusions

Our methodology enables the stabilization of the learned feature space by effectively capturing the latent feature distribution information. Experimental results demonstrate that our model considerably outperforms traditional discriminative segmentation methods across a variety of datasets from multiple modalities.