ICLR‘25 Spotlight | LeFusion：通过病变关注的扩散模型实现可控病理合成

AITIME论道

于 2025-04-10 15:13:17 发布

阅读量14

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzIzOTAxOTk3Nw==&mid=2247533246&idx=1&sn=c079aa10abd750958f6ad67743610919&chksm=e85facfac81435a642fd5810973dfb9e6ebbf2a5c6a9bf8c9eefce259daf9bba4ffb630f27e7&scene=126&sessionid=0

版权

点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

点击 阅读原文 观看作者讲解回放！

个人信息

作者：张翰韬，中国科学技术大学硕士，EPFL访问硕士

杨健程，瑞士洛桑联邦理工学院 (EPFL) 博士后，主要研究医疗AI和3D视觉。本科和博士毕业于上海交通大学。曾作为访问学者在哈佛大学和EPFL从事研究。已发表50余篇研究论文，包括Cancer Research / eBioMedicine / TMI / MedIA / CVPR / MICCAI / NeurIPS / ICLR等顶刊顶会，谷歌学术累计引用约4000次。担任MICCAI 2024/2025 和MIDL 2025的领域主席。他曾入选世界人工智能大会 (WAIC) 云帆奖、斯坦福“全球前2%顶尖科学家榜单”和Forbes 30 Under 30 Asia。他是MedMNIST的作者。个人主页：https://jiancheng-yang.com

内容简介

真实世界临床实践中的患者数据常常面临数据稀缺和长尾不平衡的问题，导致结果存在偏差或算法不公平性。本研究通过从无病变图像生成包含病变的图像分割对来解决这些挑战。以往的医学图像合成研究在分离病变信息与背景方面存在困难，导致背景质量低下且对合成输出的控制有限。受基于扩散的图像修复技术的启发，作者提出了 LeFusion，这是一种以病变为中心的扩散模型。通过重新设计扩散学习目标以专注于病变区域，作者简化了学习过程，并通过将正向扩散的背景上下文整合到逆向扩散过程中，提高了对输出的控制能力，同时保留了高保真度的背景。此外，作者还解决了病变纹理合成中的两大主要挑战：1）多峰值和2）多类别病变。作者引入了两种有效的策略：基于直方图的纹理控制和多通道分解，从而能够在复杂场景中实现高质量病变的可控生成。此外，作者还引入了病变掩膜扩散，允许对病变的大小、位置和边界进行控制，从而增加病变的多样性。在3D心脏病变MRI和肺结节CT数据集上验证表明，LeFusion 生成的数据显著提升了包括 nnUNet 和 SwinUNETR 在内的最新分割模型的性能。

论文地址：

https://arxiv.org/pdf/2403.14066

代码链接：

https://github.com/M3DV/LeFusion

Background

医学影像领域的一个重要挑战在于数据集的缺乏，特别是患病且带有标注的医学影像数据。以肺结节为例，肺结节在整个胸部轮廓中仅占很小的一部分。在训练过程中，通常提取病灶部分的图像块（patch），并将其输入相应的神经网络进行训练。然而，从全局角度来看，这些 patch 在整个胸部CT图像中的占比非常有限。以像素计量，99.8%的数据并未被有效利用。

如何充分利用这些未被利用的医学影像数据？一种简单的方法是，在正常医学影像中随机生成伪病灶。这可以通过类似于“copy-paste”的形式，将已有病灶复制到新的正常区域，从而获得带有病灶的影像。此方法不仅能够扩展可用数据集，还可以提升神经网络的训练效果，有助于医学影像分析的进一步发展。

第二个研究的动机在于医学影像中常常出现明显的长尾问题。某些疾病，尤其是直肠癌等癌症，其发病率极低，这是由其自身特性决定的。因此，在收集的数据中，不可避免地存在严重的长尾分布，即病灶数据非常稀少，而正常病例占据大部分。面对罕见病时，如何平衡数据集的分布成为一个有趣且重要的研究点。

为了应对这一挑战，本研究引入了病灶合成的主题。通过生成伪病灶，可以在数据集中创建更多患病病例，从而平衡数据分布。

Previous Methods

此前的病灶合成方法中，最简单的一种方式是手工制作（hand-craft）。手工合成的优点在于其高度可控：研究者可以在正常区域中随机选择一个区域，然后通过人工模拟一些高斯纹理和噪声，将这些手工设计的纹理直接贴合在图像中。手工方法的优点在于高度可控，能够精确选择区域并生成病灶。然而，其缺点在于难以扩展，因为每种纹理都是针对特定病例设计的，且在贴合病灶时难以保证与周围区域的一致性。

因此，在后来的其他一些研究工作中提出了使用可控扩散模型进行病灶生成。这种方法通过学习的方式，将需要生成的 mask 位置及相应纹理信息结合起来，通过可控的扩散模型实现。这种方法的优点在于具有良好的扩展性，是一种参数化的方法。然而，其缺点在于医学影像中的背景非常复杂，难以学习。医学背景与自然场景背景不同，医学背景需要符合一定的解剖结构，某些器官的位置必须严格按照序列。在这种复杂的医学背景下，如果没有任何约束，直接使用扩散模型进行学习很难控制，产生的背景往往不具备有效的医学结构。

这一问题引出了新的研究方向：在医学影像生成过程中，是否有必要对背景进行学习？在医学影像中，正常背景数据非常丰富，是否可以利用这些正常背景数据，只对病灶部分进行学习，而忽略背景部分的学习呢？

以下是之前方法的一些示例展示。通过手工合成的方法可以直接生成病灶，但其与背景的交互性较差。而采用可控扩散模型的方式时，背景的真实性和一致性也难以得到保证。左侧为心脏的示例，右侧为胸部肺结节的示例。这些图像展示了不同方法在处理病灶生成中的效果对比。

这引出了本研究的主题——Lesion-Focused（病灶聚焦）方法。Lesion-Focused方法的核心技术在于，通过修改扩散模型的学习目标，仅学习病灶区域的纹理分布。在这一方法中，通过掩码掩码作用在损失函数上，使模型仅学习病灶区域的纹理。

LeFusion

1.Lesion-Focused Diffusion

在扩散生成过程中，病灶纹理部分仍由扩散模型生成，而背景部分则采取不同的处理方式。例如，在某个时刻，原始图像通过添加对应时刻的噪声，获得时刻已加噪的图像信息。然后，保留图像的信息背景，通过掩码将背景和生成的病灶纹理结合起来。具体方法是：用1减去掩码（1-Mask）乘以背景部分，再加上病灶纹理，通过这样的形式实现背景的部分由原始图像加噪形成。因此，理论上背景部分是完全保真的，而病灶部分则由扩散模型生成纹理，从而保证了生成效果的准确性。

由于整个模型通过参数化形式构建，可以有效保证病灶与背景之间的融合，同时提高模型的泛化能力。核心技巧在于仅生成病灶部分，而不是尝试生成所有背景部分，从而实现更加精确和高效的病灶合成。

尽管仅生成病灶，而不考虑背景，看似能够解决医学影像生成问题，但实践中发现这一方法并不完全有效。有趣的是，尽管通过扩散模型生成病灶纹理，但实际影像中的病灶纹理极为复杂。例如，肺结节根据纹理可在医学上划分为多种类型，包括实心结节、毛玻璃样结节，以及较浅的实心结节等。在临床中，医生建议切除某些类型的结节，而对于其他浅表结节，则不建议进行手术。因此，区分不同纹理形状的病灶在医学上具有重要意义。

2.Controllable Pathology Synthesis

然而，单纯依靠扩散模型学习病灶纹理，仅能捕捉纹理的整体分布，难以精确生成特定类型的结节。扩散模型通过在分布中随机采样点，每次生成的结果可能不同。然而，在实际训练过程中，往往希望模型能够生成特定类型的病灶。因此，本文引入了直方图控制的方法来解决这一问题。

通过使用直方图控制，不仅可以在扩散模型中生成特定类型的病灶，还可以更好地满足医学影像生成的实际需求。直方图控制方法能够使模型在生成过程中，更加精确地采样和生成特定类型的病灶，从而提高模型在医学影像生成中的应用效果。

在训练过程中，本研究通过引入病灶的直方图来建立控制映射。这种方法的优点在于，可以通过指定直方图获得特定类型的病灶，再与背景融合在一起。对于肺结节的形状，某些研究直接使用真实形状，但这种方式可能缺乏泛化能力。因此，本研究采用生成病灶形状的方法。

3.DiffMask

在生成病灶时，本研究引入一个控制球，通过该球的大小相应控制病灶的大小。同时，控制球的位置也可以引导病灶在图像中的生成位置。具体过程中，通过在训练阶段设置控制直方图和病灶聚焦的损失函数，使模型能够在生成过程中实现精确控制。

在推理阶段，通过病灶聚焦的逆扩散方法，结合控制直方图生成目标病灶，并利用掩码融合生成的病灶和背景。在掩码生成过程中，首先通过噪声和控制球生成病灶形状，再经过迭代和后处理步骤，确保病灶在图像中的位置和形状符合预期。

这是本研究方法的总体概述。作者的方法通过解耦病灶和背景的生成，以再绘制（repaint）形式实现。该方法采用了一种“Lesion-Focused”策略，同时引入球形信息作为控制信息，有效保证了生成时病灶的定位及其大小。

Experiment

本研究选用了两个数据集进行实验，一个是肺结节数据集，另一个是心脏数据集。具体来说，选用的肺结节数据集是 LIDC-IDRI 数据集，而心脏数据集则是 EMIDEC 数据集。

在 LIDC-IDRI 数据集中，包括2624个病理区域（ROI）和3076个正常区域。训练集中包含2104个病理样本，测试集中包含520个病理样本。尽管在训练过程中通常选取ROI进行训练，本研究发现，对于病灶生成而言，使用CT的完整图像和ROI部分的效果是相同的，因为生成病灶的过程不受截取区域的影响。

在EMIDEC数据集中，包括67个病理区域。训练集中包含57个病理样本，测试集中包含10个病理样本，此外还有33个正常样本。该数据集采用核磁共振（MRI）成像。

可以发现，通过采用“Lesion Focused”策略，对下游任务的性能提升具有显著效果。左侧为肺结节的实验结果，右侧为病灶的实验结果。特别是在处理长尾分布时，Dice PMO 从36.32显著提升到43.41，体现了该方法的有效性。

在医学影像研究中，除了模型架构的改进外，更大的挑战在于数据集的训练问题。本研究方法通过生成高质量的病灶图像，显著提升了分割模型在长尾分布数据集上的表现。

这是直方图控制效果的可视化展示。图中显示，通过输入不同的直方图，可以定点采样生成特定类型的病灶。而在不加直方图控制的情况下，浅色背景病灶往往容易与背景融合，表现出较浅的状态。通过直方图控制，可以根据需求选取相应的直方图比例，从而实现多样化的数据集生成。

直方图控制的一个优点是，在生成数据集时，可以灵活调整各种病灶的比例，平衡病灶与正常背景的比例，从而提高模型的学习效果。为了评估这种生成方法的有效性，本文还对各种指标进行了对比。然而，作者认为在病灶合成领域中，下游任务如分割任务应为较有效的评估。因为虽然某些图像层次的质量指标可能显示较高分数，但其结构可能不符合预期。因此，本文采用下游任务的性能提升作为主要评估标准，在图像层次的质量评估中也取得了一定的好成绩。

本文对比了多种病灶生成方法，包括手工合成方法、之前的一些方法，以及本文提出的方法。具体进行了少样本数据集上的标注实验，比较了手工合成方法、直接控制方法和本文提出的加直方图控制的方法。

Visualization

可以观察到，手工合成方法生成的图像与背景融合显得非常突兀；而直接控制方法生成的图像背景则显得模糊且失真。相比之下，本文提出的方法能够精确控制生成效果，同时保证生成图像的纹理一致性。

以下是生成掩码的示例可视化：(a) 图示为真实的肺结节形状；(b) 图示为手工生成的掩码，手工生成方法可以指定掩码的尺寸；(c) 图示为通过本研究引导的球生成的掩码，通过调整球的大小来控制肺结节的生成尺寸，并根据球在真实图像中的位置来引导肺结节生成的位置；(d) 图示为 DiffMask 生成的方法。