Diff-UNet: A diffusion embedded network for robust 3D medical image segmentation

Abstract

Benefiting from the powerful generative capabilities of diffusion models, recent studies have utilized these models to address 2D medical image segmentation problems. However, directly extending these methods to 3D medical image segmentation slice-by-slice does not yield satisfactory results. The reason is that these approaches often ignore the inter-slice relations of 3D medical data and require significant computational costs. To overcome these challenges, we devise the first diffusion-based model (i.e., Diff-UNet) with two branches for general 3D medical image segmentation. Specifically, we devise an additional boundary-prediction branch to predict the auxiliary boundary information of the target segmentation region, which assists the diffusion-denoising branch in predicting 3D segmentation results. Furthermore, we design a Multi-granularity Boundary Aggregation (MBA) module to embed both low-level and high-level boundary features into the diffusion denoising process. Then, we propose a Monte Carlo Diffusion (MC-Diff) module to generate an uncertainty map and define an uncertainty-guided segmentation loss to improve the segmentation results of uncertain pixels. Moreover, during our diffusion inference stage, we develop a Progressive Uncertainty-driven REfinement (PURE) strategy to fuse intermediate segmentation results at each diffusion inference step. Experimental results on the three latest large-scale datasets (i.e., BraTS2023, SegRap2023, and AIIB2023) with diverse organs and modalities show that our Diff-UNet quantitatively and qualitatively outperforms state-of-the-art 3D medical segmentation methods, especially on regions with small or complex structures. Our code is available at the following link: https://github.com/ge-xing/DiffUNet.