PDM: A Plug-and-Play Perturbed Multi-path Diffusion Module for Simultaneous Medical Image Segmentation Improvement and Uncertainty Estimation
摘要
Segmentation is a key step in medical image analysis. However, previous state-of-the-art segmentation deep models are largely deterministic, i.e. outputting a class-wise binary segmentation without uncertainty estimation. Further improvement for these already deployed segmentation models with uncertainty estimation is therefore highly desirable. In this work, we proposed a simple and efficient plug-and-play module, named Perturbed Multi-path Diffusion Model (PDM), that can be directly concatenated to segmentation results from previous methods to simultaneously improve the segmentation and produce a segmentation uncertainty estimation. In this module, we first randomly perturb the segmentation input with varied morphological operations, generating multiple perturbed segmentation. Then, based on the conditional Denoising Diffusion Probabilistic Model (cDDPM), we proposed to use the noise-added perturbed segmentation as initial inputs to multiple cDDPM reverse paths. The final outputs of these paths are ensemble to produce an improved segmentation, where the uncertainty is also calculated via the pixel-wise standard deviation from these reverse results. Lastly, we design a cascade framework with densely connected averaging where PDM are embedded for further improving the performance. We collected three different segmentation datasets and demonstrated that our proposed method can consistently improve the previous methods’ segmentation while producing reasonable uncertainty maps that are potentially useful in clinics.