PhoenixPeng's blog.

Medical Diffusion model——9.FDDM

2024/01/15

一、引言

  磁共振(MR)和计算机断层扫描(CT)成像在医学诊断和治疗计划中发挥着关键和补充作用。 MR 提供卓越的软组织对比度,而 CT 通常是骨解剖扫描的首选。对于放射治疗,MR 扫描通常用于肿瘤定位和分割,而 CT 扫描提供电子密度信息用于计算放射治疗计划剂量。当前的临床实践通常将 MR 和 CT 作为单独的扫描来满足不同的需求,从而产生额外的成像时间、成像辐射剂量(来自 CT)和医疗费用。这些连续采集还会引入 MR/CT 重合失调,从而影响解剖结构的对齐和定位。获取单一模态(例如MR)并通过数值方法将其转换为另一种模态(CT)(MR-to-CT翻译)的前景可以从根本上解决上述挑战。然而,由于 MR 和 CT 扫描在图像强度和特征方面存在巨大差异,这种 MR 到 CT 的转换仍然具有挑战性。

  目前主流的图像翻译方法通常基于生成对抗网络(GAN)或变分自动编码器(VAE)。这些方法在保留解剖结构(高忠实度)方面显示出优势。然而,它们存在许多问题,包括模式崩溃、鉴别器过早收敛或分辨率下降,从而影响翻译图像的质量(低真实感)。最近,扩散模型展示了卓越的生成能力,实现了令人印象深刻的 Frechet 起始距离分数 (FID),用于衡量生成图像的真实度。代表性的扩散框架是去噪扩散概率模型(DDPM),它利用马尔可夫链蒙特卡罗(MCMC)过程逐步将高斯噪声添加到图像中,然后是映射噪声的学习反向扩散过程回到图像。然而,尽管具有潜力,扩散模型在医学图像应用中发现了重大挑战,特别是在图像转换过程中保持和维持解剖结构的完整性(低忠实度)。不断添加随机高斯噪声的前向扩散过程逐渐消除高空间频率的结构细节。这些结构信息对于医学诊断/治疗至关重要,并且反向扩散过程很难完全恢复。

  一般来说,扩散模型虽然擅长创建高度逼真的图像,但通常很难保持翻译后的解剖结构的忠实度。在不成对的成像数据集上开发扩散模型尤其困难,这些数据集缺乏良好对应的解剖结构的图像对来作为网络训练的额外线索。然而,在医学领域,未配对的数据集比配对的数据集更普遍且更容易获得。因此,非常需要一种能够从不成对的数据集(无监督学习)中学习准确的图像到图像转换的扩散模型。

  为了解决这一未满足的需求,作者提出了一种新颖的基于扩散模型的方法来实现无监督的图像到图像的翻译。该方法称为频率解耦扩散模型(FDDM)利用医学图像的解耦空间频率信息来指导保留解剖结构的图像转换。FDDM由无监督变频模块基于DDPM的扩散模型组成。无监督频率转换模块可以基于 GAN 或 VAE,执行初始 MR 到 CT 转换。基于 GAN 或 VAE 方法的优缺点,这种转换可以生成高解剖精度(忠实度)的 CT 图像,但图像质量(真实感)较低。通过变频模块输出两路,一路为高频CT信息(如傅里叶域),另一路为全频CT信息。与解剖结构边界相关的高频 CT 信息利用了基于 GAN 或 VAE 的频率转换模块在保留解剖结构方面的优势,可用于调节后续的扩散模型以保留解剖结构。另一方面,全频 CT 被馈送到对应于以下扩散模型的前向扩散过程。前向扩散充当低通滤波器,捕获整体 CT 强度和语义内容。然后,扩散模型将其输出用作反向扩散的起始步骤,以高频 CT 信息为条件,逐渐恢复 CT 图像的中频信息。与无监督变频模块生成的原始全频CT相比,条件扩散模型生成的新CT提高了真实感和图像质量,同时保留了准确的解剖结构。

  总之,本文的主要贡献是:

  1️⃣引入了一种新颖的频率解耦扩散模型(FDDM),用于无监督的 MR 到 CT 图像转换。 FDDM 在目标 CT 域中生成高度逼真的图像,同时保留源 MR 域解剖结构的完整性。 FDDM利用变频模块的结构保留能力和扩散模型的真实图像生成优势,实现了平移图像的高真实感和高忠实度。

  2️⃣理论证明扩散模型的前向扩散过程可以近似为低通滤波器,从频率转换模块提取的高频信息允许额外的解剖结构保存。此外,通过灵活的方案,可以选择性地将变频模块输出的低频和高频CT信息引入到扩散模型中,使模型进一步纠正其在最终平移的CT图像中的残余误差。

  3️⃣与现有最先进的方法(包括基于 GAN、VAE 和其他扩散模型的方法)相比,FDDM 在 MRtoCT 转换方面表现出了卓越的性能。

二、方法

  给定源 MR 图像 \(s_0\) 和目标 CT 图像\(x_0\),我们使用无监督频率转换模块将 MR 图像的高频和全频信息转换为 CT 图像的对应信息。变频模块可以是基于 GAN 的,也可以是基于 VAE 的,在结构保存方面表现出色。在 FDDM 的实现中,我们使用了 UNIT,一个基于 VAE 的框架。随后,将变换后的全频 CT 图像输入到前向扩散过程中,该过程充当低频信息提取的低通滤波器。扩散模型使用该低频信息作为反向扩散的起始步骤。反向扩散以来自频率转换模块的高频 CT 信息为条件,生成最终的平移 CT 图像。在下面的段落中,我们首先描述了扩散模型的基础知识和频率滤波器,并从理论上证明了扩散模型的前向扩散过程起到了低通滤波器的作用。然后我们描述了FDDM使用的基于具有周期一致性的VAE的MR-CT变频模块。继频率转换模块之后,我们介绍了 FDDM 中采用的独特、灵活的频率调节方案,以最大限度地提高平移图像的精度并纠正频率转换模块的残余误差。

1.用于频率解耦的扩散模型和滤波器

  接下来,在扩散模型的前向过程中,可以将 \(\textcolor{blue}{x_t}\) 表示为:

  其中\(\textcolor{blue}{z}\)代表高斯白噪声。在频谱中:

  \(\textcolor{blue}{\mathcal{F}}\)表示傅里叶变换,\(\textcolor{blue}{\mathcal{X}_t(u,v)}\)\(\textcolor{blue}{\mathcal{X}_0(u,v)}\)\(\textcolor{blue}{\mathcal{Z}(u,v)}\)是相应的频域表示。对于高斯白噪声\(\textcolor{blue}{\mathcal{Z}(u,v)}\),期望和方差为:

  高斯白噪声的自相关函数可以表示为:

  假设 \(\textcolor{blue}{δ(τ_1, τ_2)}\) 代表狄拉克\(\textcolor{blue}{δ}\)函数,高斯白噪声在所有频段上具有均匀的功率。在一般图像中,功率谱密度(PSD)与空间频率成反比:

  其中 \(\textcolor{blue}{k}\)\(\textcolor{blue}{a}\) 是缩放/修改因子。当图像中引入噪声时,噪声的 PSD 会与图像的 PSD 相结合。因此,定义步骤\(\textcolor{blue}{t}\)处的信噪比 (SNR):

  由于\(\textcolor{blue}{\alpha_t}\) 随着扩散步长单调递减,因此当扩散步长\(\textcolor{blue}{t}\)增加时,SNR 会下降。给定阈值 SNR \(\textcolor{blue}{\phi}\),超出特定步骤\(\textcolor{blue}{t}\),仅保留低频信息:

  如果我们将此步骤表示为\(\textcolor{blue}{\tilde{T}}\) ,则前向扩散反映了低通滤波器的效果:

  近似地,我们可以将反向扩散视为反向低通滤波器:

  对于 FDDM,采用高通滤波器(Sobel 算子)来提取高频信息。该算子充当离散微分算子,有效计算图像的强度梯度以提取高频分量。

  基于用作低通滤波器的正向扩散过程的输出,可以提取低频CT信息作为反向扩散的扩散模型的起始步骤。同时,Sobel滤波器可以提取高频CT信息来调节反向扩散过程以保留解剖结构。在这方面,扩散模型可以仅在CT域上进行训练,而不使用来自 MR 域的任何信息。然而,由于MR的低频和高频信息与CT的低频和高频信息不匹配,因此我们需要在仅CT扩散模型之前通过另一个频率转换模块将MR的这些信息进一步转换为CT。

2.变频模块

  由于 MR 和 CT 在全频谱上具有截然不同的信息,因此我们不能使用仅根据 CT 信息训练的扩散模型将 MR 转换为 CT。然而,我们可以首先通过频率转换模块将 MR 转换为 CT,然后从转换后的 CT 图像中提取低频和高频信息,输入到经过 CT 训练的扩散模型中,以生成微调的 CT 输出。我们基于广泛使用的具有循环一致性的 VAE 框架构建了我们的变频模块。具体来说,我们的频率转换模块有两个编码器 \(E_s\)\(E_x\)、两个解码器\(D_s\)\(D_x\) 和两个鉴别器 \(Dis_s\)\(Dis_x\),基于两个编码器和两个解码器之间共享潜在空间 L 的假设,如图 2 所示。

  该潜在空间可以从每个域独立导出,并且可以从该潜在空间重新生成 MR 和 CT 的全频图像和高频图像。具体来说,我们假设存在以下函数:

  

3.通过灵活的频率调节最大限度地减少反向扩散误差

  通过变频模块得到CT的全频图像\(\textcolor{blue}{x_0}\)和高频图像\(\textcolor{blue}{H_{x_0}}\)后,通过\(\textcolor{blue}{x_0}\)的前向扩散(低通滤波)进一步得到低频CT图像\(\textcolor{blue}{\tilde{x}_T}\)。与传统的扩散模型设置相比,FDDM设计了频率信息调节方案,在反向扩散过程中灵活地引入和去除条件信息\(\textcolor{blue}{H_{x_0}}\),这有助于校正\(\textcolor{blue}{H_{x_0}}\)的残余误差。为了通过反向扩散从噪声 \(\textcolor{blue}{x_t}\) 预测\(\textcolor{blue}{x_0}\)\(\textcolor{blue}{x_{t−1}}\)通过反向条件分布\(\textcolor{blue}{q_σ(x_{t−1}|x_t, x_0)}\) 生成,并带有中间预测的\(\textcolor{blue}{x_0}\)\(\textcolor{blue}{σ}\)表示调整更新方案的参数。基于DDIM,我们可以定义\(\textcolor{blue}{q_σ(x_{t−1}|x_t, x_0)}\) 为:

  反向扩散过程\(\textcolor{blue}{p^{(t)} _θ (x_{t−1}|x_t) }\)是一个可训练的生成过程,它利用了\(\textcolor{blue}{q_σ(x_{t−1}|x_t, x_0)}\)的知识

  根据方程 5,我们可以通过以下方程预测给定 \(\textcolor{blue}{x_t}\)\(\textcolor{blue}{x_0}\)

  其中\(\textcolor{blue}{\epsilon_\theta}\)是我们的模型,在其训练过程中,可以引入高频信息\(\textcolor{blue}{H_{x_0}}\)作为条件:

  我们通过以下变分推理目标优化 \(\textcolor{blue}{\theta}\)

  其中\(\textcolor{blue}{[H_{\bar{x}_0}]}\) 表示\(\textcolor{blue}{H_{\bar{x}_0}}\) 在训练过程中对于模型来说是可选的。具体来说,我们随机放置 Hx0 或 −1 值(无条件)。我们的反向条件分布包括是否添加高频信息,如下所示:

  我们将 \(\textcolor{blue}{\bar{e}_t}\)\(\textcolor{blue}{e_t}\) 定义为每个时间步 \(\textcolor{blue}{t}\)的反向扩散过程的误差项。这些术语衡量实际逆过程与我们模型的预测之间的差异,其中\(\textcolor{blue}{\bar{e}_t}\)考虑了额外的高频信息\(\textcolor{blue}{H_{\bar{x}_0}}\)

  然后在每一步确定最小误差项 \(\textcolor{blue}{e_{min}}\)。由于我们使用高频信息作为条件,因此我们的反向扩散可以近似视为反向低通滤波器。我们假设存在\(\textcolor{blue}{\bar{T}}\) ,当 \(\textcolor{blue}{T}\) < \(\textcolor{blue}{\bar{T}}\)时,\(\textcolor{blue}{e_t}\)和 \(\textcolor{blue}{\bar{e}_t}\) 之间的关系发生变化,如图 3 所示。

  在整个反向扩散步骤中,without/with \(\textcolor{blue}{H_{\bar{x}_0}}\) 的总误差定义为 \(\textcolor{blue}{E}\)\(\textcolor{blue}{\bar{E}}\)

  \(\textcolor{blue}{E_{min}}\)是指整个反向扩散过程中误差最小化,如下:

  推导表明,反向扩散过程中存在一个最佳步骤T,超过该步骤,放弃高频条件可以进一步减少平移图像中的误差,可以在测试时优化期间进行定制。总之,整个MR到CT的翻译过程如算法1所示。

三、实验

  在本文的研究中,采用了 SynthRAD2023 数据集1,其中包含各种神经影像扫描集合。该数据集总共包括 360 张 3D 图像,其中 180 张脑部 CT 扫描和 180 张配准的 MR 扫描来自三个不同的医院。

  实验设置,作者随机划分这些图像,留下 18 对 CT 和 MR 病例作为测试集。其余 162 个配对 MR-CT 集用于训练。每个病例都被处理为 2D 切片,从而产生 31,142 个 MR 和 31,142 个 CT 切片用于训练,以及 3,417 个 MR 和 3,417 个 CT 切片用于测试。为了模拟未配对(无监督)训练,作者重新排列切片并从训练数据集中删除配对信息。所有切片的大小均统一调整为 256 × 256 像素的分辨率。对于 MR 图像,我们应用了强度截断,删除了强度值的前 0.5%。对于 CT 图像,这些图像中的 Hounsfield 单位 (HU) 值被截断为 [-1000, 1000] 范围。然后所有图像都重新缩放为 [0, 1]。使用训练 MR 和 CT 切片来训练频率转换模块而仅使用训练 CT 切片集来训练仅 CT 扩散模型

1.与其他方法比较

  表 1 对 FDDM 与其他八种图像转换方法在 FID、SSIM 和 PSNR 方面进行了定量比较。其中,GcGAN、CycleGAN、RegGAN是基于GAN的方法; UNIT和MUNIT是基于VAE的方法; SynDiff、SDEdit 和 EGSDE 是基于扩散的方法。与 FDDM 相比,基于 GAN 和 VAE 的方法呈现出更差的 FID 分数(更高的值),显示出这些方法在生成高度逼真的域图像方面面临的挑战。相比之下,这些方法在 SSIM(一种结构相似性度量)中呈现出与 FDDM 类似的结果,符合 GAN 和 VAE 擅长保留解剖结构(忠实性)的预期,但仍被 FDDM 超越。

  FDDM 的卓越性能凸显了其通过结合变频模块的结构保持优势和扩散模型的成像真实感优势,巧妙地保持生成图像的真实感和忠实度的能力(图 4)。与依赖 CycleGAN 生成的 CT 的全频率信息来训练扩散模型的 SynDiff 方法相比,FDDM 仅使用来自基于 GAN 或基于 VAE 的 MR-CT 转换网络的选定频率信息,并填充其他信息通过在真实 CT 图像上训练的扩散模型来获取,以提高图像的真实感。另一方面,证明FDDM灵活的频率调节设计不仅有助于填补缺失的频率信息,而且进一步改善了从变频模块传递过来的频率信息。

2.扩散模型和高频调节的消融研究

  FDDM 集成了一种独特的方法,该方法使用 MR-CT 频率转换模块输出的高频和低频信息来生成所有频率的最终 CT 图像。高频和低频信息是扩散过程中关键的调节输入。为了了解这些输入的影响,我们进行了消融研究。表 2 和图 5 总结了这些消融实验的结果。

  作者首先直接测试了变频模块(无扩散)生成的全频图像的性能,如表 2 中的 FDDM w/o Diffusion 所示。没有扩散模型,FDDM 生成的结果更差,这表明 FDDM 的输出基于VAE的变频模块还可以进一步改进。对于另一项引入扩散模型但去除高频信息调节的消融研究,FID 显着改善,显示了使用扩散模型生成更真实图像的优势。然而,高频调节消融 FDDM 的 SSIM 不如完整 FDDM 模型,这表明使用高频调节保留解剖结构的好处。

3.前向扩散步骤的影响(定制低通滤波器)

  对于FDDM,虽然扩散模型是用1000个正向和反向扩散步骤进行训练的,但是可以在测试期间定制这样的步骤数,以使扩散模型接受包含不同程度的噪声/信息的输入。我们可以调整前向扩散步骤来定制传递到后续扩散模型的低频信息量,以仅保留来自变频模块的关键和高保真信号以优化输出。过度的前向扩散(低通滤波)可能会导致关键信息的显着丢失,从而影响翻译精度,如图 6 所示。而使用太少的前向步骤可能会在扩散模型中引入过多的残余误差。因此,根据表 3 中详细的测试,我们最终确定在测试时间 ̃ T 期间使用 500 个步骤是最佳的,从而实现 FID 和 SSIM 之间的平衡。

4.高频条件步骤的影响

  理论上,反向扩散存在一个最佳步数T,超过这个步数,包含来自变频模块的不完善的高频信息可能会导致更多的误差,比起去除它。表 4 中详述的实验表明,添加高频信息进行调节直至步骤 200 可实现 FID 和 SSIM 之间的最佳平衡。此外,图 7 提供了对延迟或提前停止添加高频信息的深入了解。这两种情况都会引入更多的错误,识别这个最佳步骤 T 对于最小化错误并提高最终生成图像的整体质量和忠实度至关重要。

四、总结

  在这项研究中,我们引入了频率解耦扩散模型 (FDDM),这是一种在医学成像中进行不配对 MR 到 CT 图像转换的新方法。通过结合基于 VAE 的模型在结构保存方面的优势和扩散模型在高逼真图像生成方面的优势,FDDM 实现了 MR 到 CT 转换的高保真度和真实感。 FDDM 独特的频率解耦方法使其能够使用可靠的信息来调节扩散模型以保留解剖结构。低成本的测试时间扩散步数调整允许FDDM保留最佳数量的低频信息作为输入,并在某个点停止引入高频条件,以允许扩散模型进一步纠正残余高频误差来自变频模块。 FDDM 的综合评估证明了其相对于最先进的 GAN、VAE 和扩散模型的优越性能。 FDDM 可实现高质量、解剖学上准确的 MR 到 CT 转换,促进更精确的疾病诊断、定位和治疗计划。

CATALOG
  1. 1. 一、引言
  2. 2. 二、方法
    1. 2.1. 1.用于频率解耦的扩散模型和滤波器
    2. 2.2. 2.变频模块
    3. 2.3. 3.通过灵活的频率调节最大限度地减少反向扩散误差
  3. 3. 三、实验
    1. 3.1. 1.与其他方法比较
    2. 3.2. 2.扩散模型和高频调节的消融研究
    3. 3.3. 3.前向扩散步骤的影响(定制低通滤波器)
    4. 3.4. 4.高频条件步骤的影响
  4. 4. 四、总结