[Paper Review] MedSegDiff: Medical Image Segmentation with Diffusion Probabilistic Model

Junde Wu, Fu, Rao+6 more|arXiv (Cornell University)|Nov 1, 2022

Brain Tumor Detection and Classification138 citations

TL;DR

MedSegDiff introduces a diffusion probabilistic model for general medical image segmentation, using dynamic conditional encoding and FF-Parser to improve step-wise attention and suppress high-frequency noise; it achieves SOTA on optic cup, brain tumor, and thyroid nodule segmentation across multiple modalities.

ABSTRACT

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.

Motivation & Objective

• Address general medical image segmentation with a diffusion probabilistic model (DPM).
• Develop dynamic conditioning to adapt the model at each sampling step for ambiguous medical images.
• Introduce FF-Parser to mitigate high-frequency noise during feature integration.
• Demonstrate generalization across multiple modalities (fundus, MRI, ultrasound) and organs/tat regions.
• Provide strong empirical results surpassing SOTA on three segmentation tasks.

Proposed method

• Base the model on a diffusion probabilistic framework with a UNet-style denoiser conditioned on image priors and current step segmentation.
• Introduce dynamic conditional encoding by fusing current-step segmentation features with raw image features at multiple scales via an attention-like fusion mechanism.
• Propose FF-Parser to filter high-frequency noise in the Fourier domain on skip connections during multi-scale integration.
• Train with standard DPM objective and end-to-end optimization using AdamW; accommodate multiple task datasets with different modalities.
• Evaluate with Dice and IoU metrics on optic cup (fundus), brain tumor (MRI), and thyroid nodule (ultrasound).
• Compare against SOTA CNN/ViT and diffusion-based baselines including ResUNet, BEAL, TransBTS, EnsemDiff, nnUNet, and TransUNet.

Experimental results

Research questions

• RQ1 Can a DPM be effectively applied to general medical image segmentation across multiple modalities?
• RQ2 Does dynamic conditioning improve step-wise attention and localization in ambiguous medical images?
• RQ3 Does FF-Parser reduce high-frequency noise and improve segmentation quality in diffusion-based segmentation?
• RQ4 How does MedSegDiff perform relative to SOTA methods on optic cup, brain tumor, and thyroid nodule segmentation across different modalities?

Key findings

• MedSegDiff variants achieve state-of-the-art or strong performance on three tasks across different modalities.
• Dynamic conditioning provides notable gains over vanilla DPM, especially in low-contrast or region-localization scenarios.
• FF-Parser further improves results by mitigating high-frequency noise during feature integration.
• MedSegDiff++ attains the highest reported metrics: optic cup Dice 87.5, IoU 79.1; brain tumor Dice 90.5, IoU 82.8; thyroid nodule Dice 86.6, IoU 80.2.
• Compared to EnsemDiff (a DPM-based method for brain tumors), MedSegDiff shows substantial improvements in Dice and IoU on both brain tumor and other tasks.
• The results suggest strong generalization of MedSegDiff across distinct medical image modalities.