Source-linked AI summary
MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model
Junde Wu, Rao Fu, Huihui Fang, Yu Zhang, Yehui Yang, Haoyi Xiong, Huiying Liu, Yanwu Xu
TL;DR
Medical image segmentation needs effective methods for ambiguous targets, while diffusion probabilistic models had not been applied to general medical segmentation. MedSegDiff introduces dynamic conditional encoding and FF-Parser, and it outperforms previous SOTA methods across three modalities and tasks.
Problem
Ambiguous lesions and tissues, especially in low-contrast medical images, make segmentation difficult when conditioning remains static across diffusion steps.
Method
MedSegDiff applies diffusion probabilistic modeling with dynamic conditional encoding and FF-Parser for step-wise conditioning and high-frequency-noise control.
Results
MedSegDiff outperforms previous SOTA methods across optic-cup, brain-tumor, and thyroid-nodule segmentation tasks with different image modalities.
Takeaways & Limitations
MedSegDiff provides a DPM-based scheme for general medical image segmentation and is proposed as an essential benchmark for future research.
Abstract
from arXiv · showhide
Diffusion probabilistic model (DPM) recently becomes one of the hottest topic in computer vision. Its image generation application such as Imagen, Latent Diffusion Models and Stable Diffusion have shown impressive generation capabilities, which aroused extensive discussion in the community. Many recent studies also found it is useful in many other vision tasks, like image deblurring, super-resolution and anomaly detection. Inspired by the success of DPM, we propose the first DPM based model toward general medical image segmentation tasks, which we named MedSegDiff. In order to enhance the step-wise regional attention in DPM for the medical image segmentation, we propose dynamic conditional encoding, which establishes the state-adaptive conditions for each sampling step. We further propose Feature Frequency Parser (FF-Parser), to eliminate the negative effect of high-frequency noise component in this process. We verify MedSegDiff on three medical segmentation tasks with different image modalities, which are optic cup segmentation over fundus images, brain tumor segmentation over MRI images and thyroid nodule segmentation over ultrasound images. The experimental results show that MedSegDiff outperforms state-of-the-art (SOTA) methods with considerable performance gap, indicating the generalization and effectiveness of the proposed model. Our code is released at https://github.com/WuJunde/MedSegDiff.
1 Introduction
Medical image segmentation supports diagnosis, surgical planning, and image-guided surgery, while diffusion probabilistic models have recently shown strong capabilities beyond image generation. MedSegDiff applies DPM to general medical segmentation with dynamic conditioning and FF-Parser, achieving SOTA performance across three tasks.
- Medical image segmentation partitions images into meaningful regions that support diagnosis, surgical planning, and image-guided surgery.Automatic methods may reduce manual effort and improve consistency and accuracy.
- Diffusion probabilistic models have expanded from image generation to vision tasks including super-resolution and deblurring.Recent systems demonstrated high diversity and synthesis quality.
- MedSegDiff is presented as the first DPM-based model for general medical image segmentation across different image modalities.The model targets ambiguous lesions and organs that are difficult to distinguish from background.
- Dynamic conditional encoding uses each sampling step’s image prior and current segmentation mask to enhance step-wise regional attention.The current-step mask is fused with image-prior features in a multi-scale manner.
- FF-Parser is introduced to reduce the negative effects of high-frequency noise during the adaptive integration process.The contribution list identifies FF-Parser as a component for mitigating high-frequency effects.
- MedSegDiff achieves SOTA performance on optic-cup, brain-tumor, and thyroid-nodule segmentation tasks.These tasks use fundus, MRI, and ultrasound images, respectively.
2 Method
MedSegDiff adapts diffusion-based segmentation with step-conditioned image guidance and frequency-domain feature processing. Its modified ResUNet reverses progressively noised segmentation maps while using dynamic conditioning and FF-Parser to localize regions and constrain high-frequency noise.
- Diffusion segmentation: MedSegDiff reverses a sequence of progressively noised segmentation labels, starting from Gaussian noise and recovering a clear segmentation step by step.The reverse process is parameterized by a neural network trained to undo the forward noising process.
- Dynamic Conditional Encoding: A UNet conditions each step-estimation function on raw-image features and current-step segmentation features, with the step index integrated through learned embeddings.The raw image and current segmentation map are encoded separately before their features are combined for decoding.
- Dynamic Conditional Encoding: Dynamic conditional encoding fuses same-shaped image and segmentation features using layer normalization, element-wise affinity, and attention-like enhancement of relevant regions.The fusion is applied at the middle two convolutional stages, implemented following ResNet34, to dynamically localize and calibrate segmentation.
- FF-Parser: FF-Parser transforms decoder features into Fourier space, applies a learned attentive weight map, and inverse-transforms them to constrain noise-related high-frequency components.It acts as a learnable frequency filter that globally adjusts components at specific frequencies rather than applying only spatial attention.
- Training and Architecture: The model uses a modified ResUNet with a ResNet encoder and UNet decoder, while training follows the standard diffusion-model process with sampled image-label pairs and noise.The architecture contains separate encoders for the raw image I and current segmentation state x_t.
3 Experiments
MedSegDiff is evaluated across three medical segmentation tasks and image modalities, where it outperforms competing methods and improves segmentation of ambiguous or low-contrast regions. Ablations attribute gains to dynamic conditioning and FF-Parser.
- Experimental Setup: MedSegDiff is evaluated on optic-cup, brain-tumor, and thyroid-nodule segmentation using fundus, MRI, and ultrasound images.The experiments use REFUGE-2, BraTS-2021, and DDTI datasets.
- Main Results: The comparison includes task-specific and general segmentation methods using CNN, vision-transformer, and diffusion-probabilistic architectures.Performance is evaluated with Dice score and IoU.
- Main Results: MedSegDiff outperforms all compared methods across the three tasks and modalities.Against EnsemDiff on brain-tumor segmentation, it improves 2.3% on Dice and 2.4% on IoU.
- Main Results: Visual examples show more accurate MedSegDiff segmentation, especially in ambiguous regions and low-contrast images.The authors associate this with improved localization and calibration from DPM, dynamic conditioning, and FF-Parser.
- Result Presentation: Table 1 marks the best results in bold and uses grey backgrounds to identify methods designed for particular tasks.The figure compares top-four general medical image segmentation methods across brain-tumor, optic-cup, and thyroid-nodule segmentation from top to bottom.
- Ablation Study: Dynamic conditioning improves Dice by 2.1% on optic-cup, 1.6% on brain-tumor, and 1.8% on thyroid-nodule segmentation over vanilla DPM.FF-Parser further improves performance by near 1% and achieves the best result on all three tasks.
4 Conclusion
The paper presents MedSegDiff as a DPM-based framework for general medical image segmentation, combining dynamic conditional encoding with FF-Parser. Across three modalities, it reports better performance than previous state-of-the-art methods and positions the model as a benchmark for future research.
- 4 Conclusion: MedSegDiff combines dynamic conditional encoding and FF-Parser in a DPM-based scheme for general medical image segmentation.The paper evaluates it on three segmentation tasks with different image modalities.
- 4 Conclusion: MedSegDiff outperforms previous state-of-the-art methods across the three evaluated medical image segmentation tasks.The authors describe it as the first DPM application in general medical image segmentation.
- 4 Conclusion: The authors propose MedSegDiff as an essential benchmark for future research in general medical image segmentation.