DiffBoost：通过文本引导的扩散模型增强医学图像分割
DiffBoost：通过文本引导的扩散模型增强医学图像分割
DiffBoost是一种创新的医学图像分割增强方法，通过结合文本引导和边缘信息，能够生成高质量的合成医学图像。这种方法不仅解决了医学图像分析中数据稀缺的问题，还显著提升了模型的泛化性能。本文详细介绍了DiffBoost的网络结构、训练过程和实验结果，展示了其在多个医学图像分割任务中的优越性能。
一、研究背景
人工智能（AI）和深度学习的快速发展正在推动医学图像分析领域的革命性变革。这些技术的最佳性能通常依赖于大规模、高质量的标注数据集。然而，医学领域面临着数据稀缺这一重大挑战，主要原因包括：
- 获取和注释高质量的医学图像不仅成本高昂，而且耗时、费力。
- 数据共享和收集方面的合作，可能会因隐私和伦理问题而受到限制。
- 在某些特殊疾病或极端病例中，相关数据的可获取性更加有限，进一步加剧了数据短缺问题。
当前研究通常采用数据增强的方法来扩展训练数据集，从而提升模型的泛化能力。旋转、缩放、翻转等传统方法虽然计算简单，但难以捕捉医学图像的复杂变异和高维关系。而对抗生成网络（GAN）等生成式数据增强能够学习医学成像数据中的复杂特征，并生成高逼真的样本，但由于训练不稳定、难以控制生成质量以及缺乏语义引导等问题，可能导致模型做出错误预测。
近年来，去噪扩散概率模型（DDPM）成为生成式AI领域的新热点，在高质量图像合成方面表现出色。扩散模型通过模拟随机反向扩散过程，在一系列逐步去噪的过程中，将初始噪声样本转化为逼真的医学影像，展现出较强的图像生成能力。
为了缓解数据稀缺问题，本文提出了一种文本引导的扩散模型——DiffBoost，通过结合文本信息和边缘结构，能够可控地生成高质量医学图像，提供了一种全新的数据增强策略，从而有效提升模型在医学图像分割任务中的性能。
二、方法
图 1 DiffBoost的网络结构
本文提出一种基于文本引导的扩散模型（DDPM）的数据增强方法——DiffBoost，其目的是生成可靠的合成医学图像数据，从而提升下游医学图像分割的性能。DDPM通过模拟随机反向扩散过程，将初始噪声样本逐步转换为逼真的医学图像数据。其前向过程逐步添加噪声，使数据趋近于纯噪声，而反向去噪过程则通过学习去噪函数逐步恢复原始数据。相较于使用GAN的生成方法，DDPM具有以下优势：
- 更稳定的训练过程：DDPM采用去噪得分匹配目标，优化过程更加稳定，能实现更好的收敛，而GAN由于依赖对抗训练，容易出现模式崩溃或梯度消失等问题
- 更高质量的合成医学图像：DDPM生成的样本不仅视觉上更真实，而且医学细节更清晰，符合医学影像的实际需求，而使用GAN可能会产生伪影或不符合医学规律的特征。
- 更强的采样控制能力：DDPM的采样过程是一个逐步的逆扩散过程，可以通过调整步数、噪声调度等参数，较为精确地控制生成样本的复杂度和多样性，而GAN通常无法实现如此精确的控制。
- 更强的泛化能力：DDPM采用去噪优化目标，能在数据有限地情况下进行有效训练，避免模式崩溃和过拟合，在医学图像分析任务中具有更强的鲁棒性。
图1为DiffBoost的详细网络结构，主要包括三个阶段：预训练、微调、在下游任务中使用模型。
2.1 预训练
DiffBoost使用了一个大规模的医学图像数据集RadImageNet进行训练，其中包含MRI、CT、超声等影像模态的共计135万张图像，涵盖脚踝、足部、膝盖、臀部、肩部等11个解剖区域，具有165种病理特征。
此外，本研究中还结合文本提示和边缘信息，以提高医学图像的生成质量。实验中采用“成像模态，解剖部位，病理类别”格式的文本提示（示例见图2），确保文本提示的系统性和一致性，并使用准确的医学术语，有效引导扩散模型生成符合医学结构的高质量图像。
除了文本提示外，使用Holistically-Nested Edge Detection（HED）算法精确提取到的医学图像边界作为边缘信息输入，能够确保生成的医学图像符合解剖结构。两者共同输入到扩散模型，使得生成的合成医学图像不仅在语义上符合医学标准，而且在形态结构上保持准确，为后续下游任务的微调和训练奠定基础。
图 2 文本提示示例，从左到右依次为：（CT，腹部，正常）；（磁共振，踝关节，骨质破坏)；(磁共振，大脑，正常)
2.2 微调
由于预训练可能无法完全捕捉每个医学应用特定的数据分布细节，或缺乏生成特定分割目标的能力，因此本文在特定医学任务数据集上对DiffBoost进行了微调。该步骤能使模型适应特定任务中的医学图像特征和解剖学变化，生成更相关的合成样本，从而进一步提高数据增强的有效性。
在微调阶段，模型继续结合文本提示和边缘信息作为条件输入，以确保生成医学图像的解剖结构准确。不同于预训练阶段，微调时的边缘信息来自分割掩码，以确保生成的图像在结构上与分割目标保持一致，从而提升分割任务的精准度。
2.3 下游任务的训练——图像分割
本研究中主要关注医学图像分割任务，将其作为下游应用来评估数据增强方法的有效性。经微调后的扩散模型通过添加增强文本（如增强对比度、提高分辨率等），并结合原始文本作为条件输入，生成新的合成样本用于数据增强。
这些合成样本在生成时结合了文本提示与边缘信息，确保其结构符合医学图像分布，并具备有助于分割任务的关键特征。在下游分割任务的训练过程中，合成数据与真实数据结合，使模型能够同时学习真实与增强样本，有效增强数据多样性和规模。此外，DiffBoost适配多种分割网络和损失函数，具有较强的适用性。
三、实验
3.1 经RadImageNet训练的扩散模型
文中使用平均绝对误差（MAE）、均方误差（MSE）、均方根误差（RMSE）以及多尺度结构相似性（MS-SSIM）评估生成的合成医学图像质量，并将扩散模型与传统的GAN（Pix2Pix）进行了对比。
Pix2Pix采用与扩散模型相同的边缘信息作为引导，并尝试重建原始医学图像。实验结果表明，扩散模型在所有评估指标上均优于Pix2Pix（具体见表1），证明了扩散模型在生成医学图像的真实性和多样性方面具有更高的保真度。图3展示了结合文本提示后的生成结果。
表 1 生成合成医学图像的性能评估
图 3 示例增强样本。后四列展示了生成的增强图像，前两列分别显示了原始MRI、CT和超声图像及其对应的边缘图。增强样本在保持结构完整性的同时，在强度分布方面展现出了显著变化。
3.2 下游任务上的性能
本文在多个小样本医学数据集上验证了DiffBoost的性能，并对比评估了多种传统医学数据增强方法，包括空间变换（如随机旋转、随机缩放、随机镜像等）以及强度变换（如随机对比度、随机伽马校正、随机亮度、随机噪声等）。
表2展示了DiffBoost与其他数据增强方法在不同器官分割任务上的对比。实验表明，DiffBoost在各数据集上均表现出色，尤其是在Dice系数上的显著提升。同时，DiffBoost还减少了不同数据集之间的标准差，显示出更高的特征提取一致性。此外，HD95（衡量预测分割边界与真实边界之间的平均距离）的降低进一步表明，DiffBoost生成的分割掩码在解剖结构上的形状更加精准和一致。
除了定量分析外，还通过定性分析（见图4）评估了DiffBoost的实际效果。实验发现，DiffBoost显著提升了分割的可靠性，进一步验证了其作为数据增强方法的有效性。
表 2 DiffBoost在乳腺癌、脾脏和前列腺分割任务中，与其他数据增强方法在超声、CT和MRI影像中的分割性能比较
图 4 不同数据增强方法的分割性能视觉对比
四、总结
文中提出了一种基于扩散模型的数据增强方法——DiffBoost，用于提升医学图像分割性能。DiffBoost结合了文本提示和边缘信息，能够生成具有解剖意义的逼真合成医学影像，并在多个数据集上验证了有效性。然而，该方法仍然存在一定的局限性：
- 文本输入的局限性：目前的文本输入采用三元组格式（“成像模态，解剖部位，病理类别”），未测试输入自然语言描述（如“生成一张脾脏的CT图像”）时的效果。在分割任务需要特定指令并提供更灵活的输入时，自然文本输入或许会更有效。
- 边缘信息的影响：使用ControlNet作为条件输入，被认为可能会限制生成图像的解剖学变化，降低样本的多样性。幸运的是，实验表明，通过对边缘灵敏度的层次化控制，该方法并未过度约束生成样本。但如果直接使用真实边缘（而非边缘图），可能会严重降低生成图像的多样性。
- 不同组织边界的表现：部分软组织（如肝脏、胰腺和脾脏）在HED方法的引导下可生成较为清晰的边界，但在某些情况下，软组织的边界信息仍可能模糊。因此，目前无法确定软组织的生成质量是否始终低于其他组织。
- 对比实验的局限性：研究中未对其他基于GAN方法进行对比，因为研究重点在于验证基于扩散模型的数据增强在分割任务中的可行性。未来可进一步对GAN、VAE、扩散模型这些生成式的AI算法进行更全面的对比分析，作为当前的研究补充。
基于当前研究结果，未来研究将聚焦于：
- 优化扩散模型的采样效率，加速图像生成过程，提高计算效率。
- 更有效地融入文本信息，实现对生成图像更精细的控制。
- 拓展图像条件输入模式，除了边缘信息之外，还可以从文本描述中结合更广泛的上下文，从而生成更加丰富、细致的合成医学图像。