[论文解读] Diff-UNet: A Diffusion Embedded Network for Volumetric Segmentation

Zhaohu Xing, Liang Wan+3 more|arXiv (Cornell University)|Mar 18, 2023

Radiomics and Machine Learning in Medical Imaging被引用 41

一句话总结

Diff-UNet 将扩散模型集成到三维 U-Net 用于容积医学图像分割，并使用基于步长不确定性的融合（SUF）来结合多步预测，在MRI和CT数据集上实现鲁棒结果。

ABSTRACT

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

研究动机与目标

• 促使利用去噪扩散模型实现稳健的三维医学容积分割。
• 提出一个基于扩散的端到端框架，通过标签嵌入处理多标签分割。
• 引入一个去噪模块，包含一个去噪 U-Net 和一个独立的特征编码器以提取语义信息。
• 开发一个基于步长不确定性的融合（SUF）模块，在测试阶段融合扩散步的预测。
• 在 BraTS2020、MSD Liver 与 BTCV 数据集上展示优于最先进方法的性能。

提出的方法

• 将多类分割转换为一热编码的多通道标签图，以实现基于扩散的预测。
• 使用由特征编码器和去噪 U-Net 组成的去噪模块，将体积特征与嘈杂的一热标签融合以进行 x0 预测。
• 使用结合 Dice、BCE 和 MSE 损失的综合损失进行训练。
• 在测试阶段，生成多个 DDIM 步骤，并使用基于步长不确定性的融合（SUF）加权方案融合预测，考虑步长索引和预测的不确定性。
• 不确定性通过每步的多次前向传播（类似蒙特卡洛采样）估计，并用于对最终分割输出中的逐步预测进行加权。

实验结果

研究问题

• RQ1 扩散模型能否有效嵌入到三维 U-Net 架构中，以执行多标签医学容积分割？
• RQ2 基于步长不确定性的融合机制是否提高了扩散基础分割在具有挑战性数据集上的鲁棒性？
• RQ3 与最先进方法相比，Diff-UNet 在多器官 CT 和多模态 MRI 分割基准测试中的表现如何？

主要发现

• Diff-UNet 在 BraTS2020 上取得更高的平均 Dice 分数（平均 85.35%），且在 TC/WT/ET 区域优于若干 SOTA 方法。
• 在 MSD Liver 上，Diff-UNet 在肝脏上取得 95.72% Dice，HD95 为 0.222，在肝脏肿瘤上取得 51.65% Dice，HD95 为 17.280，跨区域平均 Dice 为 73.69，HD95 为 8.751。
• 在 BTCV 上，Diff-UNet 以 83.75 的最佳平均 Dice 和 8.115 的 HD95 在八个器官中表现最佳。
• 消融研究显示，加入独立的特征编码器（FE）和 SUF 模块可提升性能，且在不确定性采样的 S=4 时获得最佳结果。
• 定性结果表明，Diff-UNet 产生的分割更精准，特别是对小目标，相较于竞争方法。