3d unet 医学图像分割（医学图像分割模型）

ImgX-DiffSeg：基于DDPM 的3D 医学图像分割

前言

本文继续解读医学图像扩散系列。之前，我们分别介绍了扩散在自监督和监督分割中的应用。关联：

ICLR 2023：基于扩散对抗表示学习的血管分割

MedSegDiff：基于扩散概率模型的医学图像分割

这次的文章《在3D多类分割中对齐训练策略与扩散模型的评估的重要性》并不是一个新的扩散应用，而是对训练和推理策略的优化，并适应了3D医学图像分割任务。参考链接在文章最后。

目前存在的问题

目前，扩散模型的架构训练和推理非常耗时。

在某些分割任务中，不确定噪声推断分割图的扩散模型预测还是分割图的直接预测更好。

该模型过于依赖先前时间步骤的信息。

ImgX-DiffSeg 架构

概述

DDPM是一种可用于图像去噪和分割的生成模型。它的工作原理是模拟干净图像的概率分布，然后向图像添加噪声以产生噪声版本。相反，该模型尝试通过去除添加的噪声来对图像进行去噪。在图像分割的情况下，模型会生成分割掩模，该掩模可以根据输入图像的特征将图像划分为不同的区域。更详细的内容，建议阅读先决条件文章（强烈推荐）。

对于ImgX-DiffSeg来说，整体流程如下图所示。首先，该架构预测分割掩模而不是采样噪声，并通过Dice Loss 直接优化。这意味着ImgX-DiffSeg 可以直接预测图像的分割图，而不是生成噪声并使用它来推断分割。其次，回收前一时间步的预测掩模以生成（噪声损坏的掩模）噪声损坏的掩模。这有助于减少信息泄漏，当模型过于依赖先前时间步骤的信息时，就会发生信息泄漏。最后，训练的扩散过程减少到五个步骤，与推理过程相同。扩散过程是一种平滑图像中噪声的方法，减少步骤数有助于提高效率。

请添加图片描述

具有方差表重采样的DDPM

训练过程的公式与DDPM基本一致。以下公式分别表示逆向过程中的预测噪声和原始图像：

在推理过程中，DDPM 中的生成过程从正常噪声开始，由变量xT 表示。该初始噪声是从平均值为0、方差为1 的正态分布中采样的。在生成过程的每个步骤中，使用预测平均值 对变量xtk-1 进行采样。下标k-1表示前一个时间步。这意味着每一步的x 值取决于上一步的x 值和分布的预测平均值。

以上流程与DDPM类似，不再赘述。重点关注Variance Schedule Resampling的过程，即如何将训练的扩散过程减少到五个步骤。可以理解为对方差值子序列进行采样的过程。给定方差表{t} Tt=1，子序列{k} Kk=1可以用{tk} Kk=1来采样。简而言之，在训练或推理过程中，给出一系列方差值，并对这些值的子序列进行采样。子序列中的值是根据先前的值和重新计算的值计算的。目标是通过在训练或推理过程中调整方差值来优化模型的性能。如果在图像去噪任务中对方差进行重采样，肯定会影响结果，但在分割任务中已经验证是有效的。

分割的扩散模型

前面的部分是DDPM的方差重采样，不涉及图像分割过程。为了优化分割任务，可以使用基于时间步长的预测噪声和采样噪声之间的L2 损失来训练ImgX-DiffSeg。此外，ImgX-DiffSeg 计算预测掩模和黄金标准之间定义的特定分割损失，例如Dice 损失或CE 损失。

在训练过程中，现有方法通过对噪声和黄金标准进行插值来对噪声掩模进行采样，这会导致数据信息的丢失。为了解决这个问题，模型的前一个时间步的预测被回收以取代黄金标准。使用以下等式计算回收的噪声掩模。

其中x0theta 是使用黄金标准计算的前一个时间步的预测分割掩模，xt 和xt+1 是两个独立的采样噪声。对xt+1 应用梯度停止，以防止通过回收的噪声掩模进行反向传播。 t 是一个超参数。第一个方程使用先前的预测和当前噪声计算xt，而第二个方程使用黄金标准和下一个噪声计算xt+1。

实验

实验分别基于MRI 和CT 图像数据集。值得注意的是，ImgX-DiffSeg 在3D 数据集上的性能优于2D 数据集，如下表所示。

下图是非扩散分割模型和扩散概率模型的直观对比，其中t代表时间步，向后扩散直到第一个时间步效果最好。

请添加图片描述

下表显示了四种消融实验，包括预测噪声推断分割图和直接预测分割图的比较；损失函数的比较；比较是否回收前一个时间步预测的mask；训练过程中时间步数的比较。 5步的效果优于1000步，说明方差表重采样是有效的。

请添加图片描述

总结

ImgX-DiffSeg 是第一个用于3D 图像多类分割的DDPM 模型。与现有的基于扩散的方法相比，该模型显着提高了性能，但并不比普通的非扩散分割模型更好。这也是值得的。进一步改进。本文的代码现已开源。我尝试了一下，发现训练收敛速度确实很快，可以作为一个很好的基准。

参考

https://github.com/mathpluscode/ImgX-DiffSeg

https://arxiv.org/abs/2303.06040

审稿人：李茜