一种新颖的扩散模型双重采样策略,DDPM + DDIM ~
论文:Diffusion Models for Counterfactual Generation and Anomaly Detection in Brain Images
0、摘要
病理区域的分割掩模在许多医学应用中很有用,例如脑肿瘤和中风管理。此外,疾病图像的健康反事实可以用来增强放射科医生的培训文件,并提高分割模型的可解释性。)
本文提出了一种弱监督方法来生成一个疾病图像的健康版本,然后使用它来获得像素异常图。为此,本文首先使用 ACAT 生成一个大致覆盖病理区域的显著性图,然后,本文提出了一种技术,允许对这些区域进行有针对性的修改,同时保留图像的其余部分。
具体而言,本文使用在健康样本上训练的扩散模型,并在采样过程的每一步结合去噪扩散概率模型(DDPM)和去噪扩散隐式模型(DDIM)。DDPM 用于修改显著性图中受病变影响的区域,而 DDIM 则确保在这些区域之外重建正常的解剖结构。这两部分在每一步都进行融合,以确保生成的样本具有一致的外观,并在编辑和未编辑部分之间实现无缝过渡。
当本文方法应用于健康样本时,输入图像在没有显著修改的情况下被重建。本文在脑病变分割任务中与替代的弱监督方法进行了比较,在所考虑的模型中获得了最高的平均 Dice 和 IoU 分数。
Figure 1 | 首先通过 DDIMs 的反向采样技术将异常图像 x 0 x_0 x0 转换为其噪声版本 x K x_K xK。随后,利用 DDPM 采样来修改通过 ACAT 生成的显著图识别出的病理性区域,旨在根据周围区域的上下文信息恢复正常的解剖结构:与此同时,图像中不含任何病理元素的区域则通过 DDIM 采样恢复至其原始外观。在整个采样过程中,这两个组件被融合在一起,以确保编辑部分与未编辑部分之间过渡自然且逼真,最终生成的图像 x ^ 0 \hat x_0 x^0 在视觉上呈现出连贯且自然的效果;
1、引言
1.1、当前挑战
(1)自动编码器(Autoencoders)并不总能生成清晰的图像,也不能保证正确地映射到健康版本;生成对抗网络(GANs)的训练有时不稳定,依赖于许多超参数,并且会产生较差的样本;
(2)已有研究使用扩散模型和分类器引导来恢复正常的解剖结构。然而,用于指导采样过程的梯度需要从未经噪声处理的样本训练的分类器计算得出。由于在医学成像中,样本的类别通常由微小细节决定,而这些细节可能在仅几次噪声处理后就丢失了,因此这种分类器经常产生不可靠的预测。因此,采用这种方法时无法保证保留样本的原始结构,许多正常组织的细节可能会被修改;
(3)尽管对抗性反事实注意力(Adversarial Counterfactual Attention,ACAT)侧重于生成反事实图像,但其主要优势在于能够准确识别病理区域,这些区域随后被用于分类流程中。然而,它在生成可信的反事实示例方面存在不足;
图 2 展示了这一现象,可以看到 ACAT 能够生成一个大致识别出病理区域的显著性图(下面那行)。然而,在反事实示例中,病变仍然可见(上面那行)。相比之下,本文的方法不仅优化了显著性图,还生成了一个病理完全被消除的反事实图像。
Figure 2 | IST-3 的原始图像(a)以及相应的健康反事实图像(第二行)和对应的异常图(底行),这些异常图是通过去噪自编码器(DenoisingAE,b)、f-AnoGAN(c)、AnoDDPM(d)、AutoDDPM(e)、分类器引导(f)、无分类器引导(g)、ACAT(h)和 Dif-fuse(i)获得的:ACAT 生成了一个合理的异常图,但无法完全移除病变。Dif-fuse 在优化 ACAT 获得的异常图的同时,还创建了一个可信的反事实示例。其他方法则引入了伪影和/或对病理区域的识别不够准确;
1.2、本文贡献
利用通过 ACAT 获得的显著性图来指导扩散模型的图像生成过程。首先,在健康样本上训练一个去噪扩散概率模型(DDPM),并结合 DDPM 和 DDIM 采样,以去除图像中的病理性区域:
(1)本文提出了一种新颖的扩散模型双重采样策略,这种策略无需病变注释,即可对由分割掩码识别出的感兴趣区域(ROIs)进行修复,同时保留图像的其余部分;
本文的创新之处在于在每个时间步混合两个组件的方法,这使得编辑部分和未编辑部分之间实现了平滑的融合。这不仅能够生成医学图像的真实反事实示例,还能生成病理区域的异常图;
(2)在 WMH 和 BraTS 2021 数据集上,本文的方法与现有的弱监督医学图像分割方法进行比较,结果显示本文方法在两个数据集中所考虑的方法中均实现了最高的平均 Dice 和 IoU 分数;
本文方法在图像质量方面也具有可比性,这通过在 IST-3 上使用核初始距离(Kernel Inception Distance,KID)来衡量,与无约束(无掩码)的扩散采样方法相比,本文方法具有更准确的异常分割这一额外优势;
2、方法
本文的采样方法不仅能够生成高度逼真的反事实图像,还能优化第一步中通过 ACAT 获得的初始显著性图。这是因为选定的区域可能不会被扩散模型完全修改,从而使得初始显著性图中识别出的健康解剖特征得以保留。方法总览如 图 1 所示。
2.1、Diffusion Models
原文略,可参考:【Diffusion综述】医学图像分析中的扩散模型(一)中 2.2 节;
2.2、Dif-fuse
本文使用了在健康样本上训练的 DDPM,并通过 ACAT 从反事实示例中获取显著图,本文选择 ACAT,因为它在识别脑部和肺部 CT 扫描中的病理性区域方面表现出色。然而,原则上,显著图也可以通过其他任何方法生成。
Dif-fuse 算法流程:
给定一个病变图像 x 0 x_0 x0,首先选择一个噪声量 K ∈ [ 0 , T ] K∈[0,T] K∈[0,T],利用逆向 DDIM 采样方案,将图像映射到其噪声版本 x K x_K xK:
随后,使用一个尺寸为 5 × 5 5×5 5×5 的高斯核对显著性图进行平滑处理,以获得一个更加均匀且孤立像素更少的掩码 m m m。接着,通过 DDPM 采样对掩码内的病变区域进行编辑。
由于扩散模型是在正常样本上训练的,这些区域被映射为健康外观。其余解剖结构需要保留,因此采用 DDIM 采样法处理掩模外的区域。为了获得一致的结果,将掩模部分与图像其余部分混合:
其中 ⊙ ⊙ ⊙ 是哈达玛积。这样,编辑过程集中在显著图捕捉到的部分,防止对扫描结构特征进行随意更改。实际上,DDIM 采样保证了不需要编辑的部分能够重建。此外,病理部分的修改由 DDPM 在考虑周围解剖背景的情况下进行。
当使用公式(9) 计算 x ^ t − 1 \hat x_{t−1} x^t−1 时,两个组件的总和可能不会产生完全一致的结果。然而,这种不一致性会在下一个扩散步骤中得到解决,该步骤能更好地融合这两个组件。
如果只是用 DDPM 计算 x ^ 0 \hat x_0 x^0,然后在采样过程结束时才应用掩模,情况就不会如此。图 3 展示了这种效果,可以看到,仅在采样过程(b)结束时应用掩模生成的正常图像,表现出一些伪影,并且编辑区域和未编辑区域之间缺乏平滑过渡。
Figure 3 | IST-3 的输入图像(a)以及仅在采样过程结束时应用掩码生成的正常图像(b):可以观察到(b)中存在一些伪影,并且编辑部分与未编辑部分之间没有平滑的过渡;
通过这种方式,能够获得给定病理图像的正常版本。为了获得异常图,首先计算原始图像与生成图像之间的差异,然后对得到的差异图依次应用腐蚀和膨胀操作,其中腐蚀和膨胀均使用 5 × 5 5×5 5×5 的核,以去除噪声,最后再次进行膨胀和腐蚀操作,同样使用 5 × 5 5×5 5×5 的核,以填补差异图中的小孔洞。
2.3、训练细节
(1)扩散模型 60000 次迭代,batch size=10;
(2)AdamW 优化器,学习率 1e-4, β 1 = 0.9 , β 2 = 0.999 β_1 = 0.9, β_2 = 0.999 β1=0.9,β2=0.999,weight decay=0.05;
(3)Improved DDPM 采用的损失;
(4)EMA:0.99,线性噪声调度: [ 1 0 − 4 , 0.02 ] [10^{-4}, 0.02] [10−4,0.02];
(5)1000 个采样步,UNet 通道为128;
3、实验与结果
3.1、数据集
(1)IST-3 数据集(CT):3035 名出现中风症状的患者的大脑影像数据;本文共考虑了5681 次扫描,其中 46.31% 被分类为阴性(无病变),其余扫描结果为阳性。对每次扫描的 11 个切片进行了处理,并将每个切片调整至 256×256。
(2)BraTS 2021 数据集(MRI):1251 名患者的扫描,每个扫描有 155 层,本文移除最上面和最下面的 25 层,利用 0 填充到 256×256,最终得到 131,164 个切片,其中 79,113 个为阳性。
(3)WMH 数据集(MRI):白质高信号强度分割挑战,使用测试集中 110 次扫描,对每个切片中心裁剪并调整大小为 256×256。
由于 IST-3 数据集中没有病变的注释,利用该数据集来评估生成图像的质量,而不是分割精度。另一方面,对于 BraTS 2021 和 WMH 数据集,能够获取病变的注释,这能够对创建的异常图进行定量分析。IST-3 和 BraTS 2021 被划分为训练集、验证集和测试集,比例为 70%、15%和15%。在 WMH 上,本文在没有进一步微调的情况下评估在 BraTS 2021 上训练的模型,以测试它们的跨领域泛化能力。
3.2、反事实示例
Figure 4 | 原始图像来自 BraTS 2021,包含真实分割掩模(a)和健康图像(上排),以及相应的异常图(下排),这些图像是通过 DenoisingAE(b)、f-AnoGan(c)、AnoDDPM(d)、AutoDDPM(e)、分类器引导(f)、无分类器引导(g)和 Dif-fuse(h)获得的:f-Ano GAN在生成可信的反事实方面表现不佳,而其他方法则产生了更高质量的结果。然而,DenoisingAE、AnoDDPM 和 AutoDDPM 并不能完全去除损伤,而使用 CG 和 CFG 生成的反事实则表现出一些伪影;
Figure 6 | 来自 BraTS 2021 的健康输入图像(a)、通过 Dif-fuse 生成的图像(b)以及异常图(c):可以观察到,本文的方法能够很好地重建健康样本;
3.3、超参数
Figure 5 | 在验证数据集上,使用不同的阈值百分位数来二值化显著性图以及不同的噪声量 K K K 所获得的 Dice 分数: 在 K = 500 K=500 K=500 以及使用显著性图中第 90 百分位的像素时,能够获得最佳结果;
Figure 7 | 来自 BraTS 2021 的原始图像及其真实的分割掩码(a),以及通过 Dif-fuse 生成的健康图像(顶部行)和对应的异常图(底部行),其中分别使用了 250(b)、500(c)和 750(d)步的加噪过程: 当使用较低的噪声量时,病理区域未能被完全移除;而当噪声水平过高时,则可能会引入显著的伪影;
3.4、定量分析
Table 1 | 在 BraTS 2021 和 WMH 测试数据上的 Dice 和 IoU 分数,以及在 IST-3 上的 KID(核初始距离),这些指标均是基于三次运行的平均值(并附上了标准误差):Dif-fuse 在 BraTS 2021 和 WMH 上均实现了最佳的异常分割性能。而采用 DDIM 反演后接 DDPM 采样的消融实验在 IST-3 上获得了最佳的 KID 值;
有点意思,朴素又实惠 ٩(๑•̀ω•́๑)۶