HiDiff: Hybrid Diffusion Framework for Medical Image Segmentation

Abstract

Medical image segmentation has been significantly advanced with the rapid development of deep learning (DL) techniques. Existing DL-based segmentation models are typically discriminative; i.e., they aim to learn a mapping from the input image to segmentation masks. However, these discriminative methods neglect the underlying data distribution and intrinsic class characteristics, suffering from unstable feature space. In this work, we propose to complement discriminative segmentation methods with the knowledge of underlying data distribution from generative models. To that end, we propose a novel hybrid diffusion framework for medical image segmentation, termed HiDiff, which can synergize the strengths of existing discriminative segmentation models and new generative diffusion models. HiDiff comprises two key components: discriminative segmentor and diffusion refiner. First, we utilize any conventional trained segmentation models as discriminative segmentor, which can provide a segmentation mask prior for diffusion refiner. Second, we propose a novel binary Bernoulli diffusion model (BBDM) as the diffusion refiner, which can effectively, efficiently, and interactively refine the segmentation mask by modeling the underlying data distribution. Third, we train the segmentor and BBDM in an alternate-collaborative manner to mutually boost each other. Extensive experimental results on abdomen organ, brain tumor, polyps, and retinal vessels segmentation datasets, covering four widely-used modalities, demonstrate the superior performance of HiDiff over existing medical segmentation algorithms, including the state-of-the-art transformer- and diffusion-based ones. In addition, HiDiff excels at segmenting small objects and generalizing to new datasets. Source codes are made available at https://github.com/takimailto/HiDiff.

Accepted by IEEE Transactions on Medical Imaging 2024

References in corpus (10)

• TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation
• A large annotated medical image dataset for the development and evaluation of segmentation algorithms
• Polyp-PVT: Polyp Segmentation with Pyramid Vision Transformers
• Using DUCK-Net for Polyp Image Segmentation
• CoreDiff: Contextual Error-Modulated Generalized Diffusion Model for Low-Dose CT Denoising and Generalization
• XNOR-Net++: Improved Binary Neural Networks
• BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation
• Multi-scale Hierarchical Vision Transformer with Cascaded Attention Decoding for Medical Image Segmentation
• ASCON: Anatomy-aware Supervised Contrastive Learning Framework for Low-dose CT Denoising
• Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation

Cited by in corpus (2)

• Dual Interaction Network with Cross-Image Attention for Medical Image Segmentation
• Resource-Aware Evolutionary Neural Architecture Search for Cardiac MRI Segmentation