PaperlayerView exact revision on arXiv

Source-linked AI summary

Diff-UNet: A Diffusion Embedded Network for Volumetric Segmentation

Zhaohu Xing, Liang Wan, Huazhu Fu, Guang Yang, Lei Zhu

arXiv:2303.10326v1eess.IVcs.CV

TL;DR

Medical volumetric segmentation requires identifying lesion areas in high-dimensional images, while prior diffusion segmentation methods are limited to 2D and cannot directly generate multi-label results. Diff-UNet embeds diffusion in a 3D segmentation network and fuses predictions across inference steps. The paper reports superiority over state-of-the-art methods across three datasets.

  • Problem

    Prior diffusion segmentation methods are limited to 2D, and conventional diffusion models cannot directly generate multi-label segmentations for volumetric medical images.

    Source: §1 · p. 2

  • Method

    Diff-UNet uses one-hot label embedding, a Denoising Module with fused volumetric features, and Step-Uncertainty based Fusion to combine iterative predictions.

    Source: §1 · p. 2, §2.1 · p. 3, §2.2 · p. 4

  • Results

    Diff-UNet is reported to outperform state-of-the-art methods across BraTS2020, BTCV, and MSD Liver; on BTCV it has the best averaged Dice and HD95 scores, 83.75% and 8.115.

    Source: §1 · p. 2, §3.2 · p. 7

  • Takeaways & Limitations

    The reported experiments support the effectiveness of diffusion models for 3D medical image segmentation.

    Source: §4 · p. 9

Abstract

from arXiv · show

In recent years, Denoising Diffusion Models have demonstrated remarkable success in generating semantically valuable pixel-wise representations for image generative modeling. In this study, we propose a novel end-to-end framework, called Diff-UNet, for medical volumetric segmentation. Our approach integrates the diffusion model into a standard U-shaped architecture to extract semantic information from the input volume effectively, resulting in excellent pixel-level representations for medical volumetric segmentation. To enhance the robustness of the diffusion model's prediction results, we also introduce a Step-Uncertainty based Fusion (SUF) module during inference to combine the outputs of the diffusion models at each step. We evaluate our method on three datasets, including multimodal brain tumors in MRI, liver tumors, and multi-organ CT volumes, and demonstrate that Diff-UNet outperforms other state-of-the-art methods significantly. Our experimental results also indicate the universality and effectiveness of the proposed model. The proposed framework has the potential to facilitate the accurate diagnosis and treatment of medical conditions by enabling more precise segmentation of anatomical structures. The codes of Diff-UNet are available at https://github.com/ge-xing/Diff-UNet

1 Introduction

Medical volumetric segmentation identifies lesion areas pixel by pixel, and more accurate results can assist doctors in diagnosis. The introduction motivates Diff-UNet by limits in existing convolutional, transformer, and diffusion approaches, then presents a 3D diffusion-based method for multi-target segmentation.

1 Introduction

  • Medical volumetric segmentation identifies lesion areas pixel by pixel in high-dimensional images, and more accurate results can provide information that assists doctors in diagnosis.Source: §1 · p. 1
  • Convolutional methods may not extract global features effectively, while transformer-based methods face computational limits on multi-scale feature extraction.Source: §1 · p. 2
  • Prior diffusion segmentation methods are limited to 2D, and conventional diffusion models cannot directly generate multi-label segmentations.Source: §1 · p. 2
  • Diff-UNet is proposed as an end-to-end 3D diffusion method for multi-target segmentation, with experiments reported to outperform state-of-the-art approaches across three datasets.The method uses one-hot label embedding, a Denoising Module, and Step-Uncertainty based Fusion (SUF).Source: §1 · p. 2

2 Method

Diff-UNet treats segmentation as denoising noisy label maps conditioned on volumetric images. Its method embeds labels for multi-target prediction, uses a denoising network with fused image features, and combines iterative predictions using uncertainty and step information.

2.1 Label Embedding

  • One-hot encoding converts the single-channel label map into multi-channel labels, which are progressively noised; the Denoising Module then predicts the clear label map from the noisy labels and raw volume.Source: §2.1 · pp. 3–4

2.2 Denoising Module

  • The Denoising Module concatenates the volume and noisy label with timestep information, while a separate Feature Encoder extracts image features for fusion with the Denoising-UNet features.Corresponding multi-scale features are summed before the Denoising-UNet decoder predicts the label map.Source: §2.2 · p. 4
  • Diff-UNet directly predicts the clean label map and combines Dice, BCE, and MSE losses for training.Source: §2.2 · p. 4

2.3 Step-Uncertainty based Fusion

  • At test time, DDIM iterates through prediction steps, and multiple predictions per step estimate uncertainty; random initialization supplies randomness without activating dropout.Source: §2.3 · pp. 4–5
  • SUF weights each step’s prediction using both step number and uncertainty, then combines the weighted predictions into the final output.Source: §2.3 · p. 5

3 Experiments

Diff-UNet is evaluated on BraTS2020, MSD Liver, and BTCV, with comparisons against state-of-the-art methods and ablations of its modules. It reports strong segmentation results across the datasets, while ablations support contributions from feature encoding and step fusion.

3.2 Comparison with SOTA Methods

  • 85.35% average Dice on BraTS2020, 1.97% above ModelsGenesis, while Diff-UNet ranks third in average HD95 at 3.3898.Its Dice score leads across WT, TC, and ET; the two higher average HD95 scores are 3.2499 and 3.0961.Source: Table 1 · p. 6, §3.2 · p. 6
  • On MSD Liver, Diff-UNet reports average Dice 73.69% and HD95 8.751, with liver Dice/HD95 of 95.72%/0.222 and liver-tumor Dice/HD95 of 51.65%/17.280.Source: Table 2 · p. 6, §3.2 · p. 7
  • 83.75% average Dice and 8.115 average HD95 on BTCV; Diff-UNet leads average scores and ranks first in Dice for six organs.It ranks second for Spleen and third for Stomach; comparison methods were not run on BTCV, so qualitative results are not shown.Source: Table 3 · p. 7, §3.2 · p. 7
  • Visual comparisons on BraTS2020 and MSD Liver show more accurate results, especially for tiny targets, while other methods sometimes miss targets or include non-target regions.The comparisons include SwinUNETR, which includes non-target regions in the illustrated example.Source: §3.2 · p. 7, Figure 2 · p. 8

3.3 Ablation study

  • Ablations on BraTS2020 report higher average Dice with the Feature Encoder, simple fusion, and then the proposed step-uncertainty weighting than with the preceding variants.The paper attributes the Feature Encoder gain to added image information and reports further gains from fusing DDIM-step predictions and weighting them differently.Source: Table 4 · p. 7, §3.3 · p. 8
  • S = 4 gives the best average Dice across the three BraTS2020 regions among the tested uncertainty-estimation settings.At S = 4, Dice is 92.23% on WT, 76.87% on TC, and 85.35% on ET.Source: Table 5 · p. 7, §3.3 · p. 8

4 Conclusion

Diff-UNet is presented as the first diffusion-based 3D medical image segmentation method, modeling segmentation as discrete data generation. Experiments on different benchmark datasets demonstrate superiority over state-of-the-art approaches.

  • Diff-UNet is presented as the first diffusion-based 3D medical image segmentation method, modeling segmentation as a discrete data generation task.Source: §4 · p. 9
  • Experiments on different benchmark datasets demonstrate Diff-UNet’s superiority over state-of-the-art approaches.Source: §4 · p. 9
  • The authors state that Diff-UNet has potential to facilitate more precise diagnosis and treatment and improve patient outcomes.Source: §4 · p. 9

Related papers

Loading 2303.10326v1…