利用合成数据生成解决医学成像的局限性
利用合成数据生成解决医学成像的局限性
医学成像中的合成数据提供了许多好处,包括在真实数据有限的情况下,以多样化和逼真的图像增强数据集的能力,从而减少与注释真实图像相关联的成本和人工。此外,合成数据还为使用敏感患者数据提供了一种合乎道德的替代方案,有助于在不损害患者隐私的情况下进行教育和培训。
这篇文章介绍了MAISI,一个NVIDIA AI Foundation模型用于3D计算机断层扫描(CT)图像生成。MAISI的首要目标是通过提供一种可靠高效的方法来生成高质量合成图像,以用于各种研究和临床应用,从而彻底改变医学成像领域。通过克服数据稀缺和隐私问题的挑战,MAISI旨在提高医学成像数据的可访问性和可用性。
该模型可以生成高分辨率合成CT图像和相应的分割掩模,最多127个解剖类别(包括骨骼、器官和肿瘤),同时实现512×512×512的标志性体素尺寸和1.0×1.0×1.0 mm³的间距。关键应用包括数据增强,即生成真实世界的医学成像数据,以补充受到隐私问题或罕见的数据集。
概述
NVIDIA的DLMED研究团队专注于3D医学图像生成建模中的高分辨率、详细背景,这种方法不仅丰富了数据集,而且提高了医学成像领域中其他机器学习模型的性能。另一个主要应用程序是节省注释工作,基于用户定义的类(图像、标签)生成对简化了创建带有注释的合成医学图像的过程,为收集和注释真实医学数据的劳动密集型任务提供了一种经济高效的替代方案。
此外,MAISI模型还解决了道德数据使用的问题,为使用敏感的患者数据提供了一种负责任的替代方案,因为生成的图像与真实的个人不符,这种能力对于生成用于教育目的的各种医学图像是非常宝贵的,帮助受训人员和医学生在不必访问机密患者记录的情况下进行诊断。
基础压缩网络
为了生成高分辨率的3D图像,研究团队训练了一个基础压缩模型,该模型旨在有效地将CT和磁共振成像(MRI)数据压缩到压缩的特征空间中。这种变分自动编码器(VAE)模型接受CT或MRI图像作为输入,并产生特征表示输出,该输出作为后续潜在扩散模型的基础输入。该模型的训练方案涵盖了来自不同解剖区域的大量CT和MRI图像,这些图像具有不同的体素间距。
这种广泛的训练赋予了模型强大的适应性,使其能够应用于不同的数据集,而不需要额外的微调。同时,一個精心训练的解码器模型可以从生成的特征集中准确重建高分辨率图像。
基础扩散网络
潜在扩散模型(LDMs)已成为生成机器学习中的一个强大工具,尤其是用于合成3D医学图像。这些模型通过从潜在空间内的随机分布中迭代地去除噪声来发挥作用,该过程有效地使LDM能够学习训练数据的底层数据分布,然后生成新颖的高保真样本。
在3D医学成像领域,LDM在生成解剖准确和多样化的图像方面具有巨大的前景,因为通过学习数据分布,该模型可以生成反映真实世界变化的合成图像。
我们的LDM是使用大规模、高分辨率CT数据集进行训练的,我们还引入了基于身体区域的条件作为额外的特征嵌入,这些区域包括头部、胸部、腹部和下半身。在推断阶段,用户可以指定他们希望为其生成CT图像的身体区域,生成的CT图像的两个具体示例如图1所示。
图1:具有不同区域输入的生成图像的示例
ControlNet支持额外的条件控制
ControlNet是一个支持各种空间上下文的框架,作为Stable Diffusion等扩散模型的附加条件,以便在文本到图像的扩散模型中添加条件控制。它是在论文《将条件控制添加到文本到图像的扩散模型中》中介绍的。有了ControlNet,用户可以更好地控制生成过程,输出可以使用不同的空间上下文进行定制,如深度图、分割图、涂鸦、关键点等。
具体而言,研究团队利用ControlNet将包括127个解剖结构在内的器官分割图作为基础扩散模型的额外条件,以促进CT图像的生成。图2显示了一个典型的生成CT图像及其相应的分割条件。
图2:典型生成的CT图像及其相应分割条件的示例
这是通过使用“零卷积”层连接可训练副本和锁定副本来实现的。零卷积层使模型能够保留预训练的基础扩散模型已经学习到的语义,同时使可训练副本能够学习任务所需的特定空间条件。
绩效评估
我们的团队使用多个数据集对基础扩散模型和ControlNet进行了全面评估。这确保了对许多不同身体区域的广泛覆盖。
图像质量
最初,我们使用提供的模型权重,通过比较模型生成的图像与其他基线方法生成的图像,来评估图像的质量。我们使用了表1所示的胸部CT图像生成和实际胸部CT数据集。
根据Fréchet起始距离(FID)评分,我们的方法在性能上比以前的方法表现出了优越。此外,我们生成的图像在外观上与真实图像非常相似。
FID (Average) ↓ | MSD Task 06* | LIDC-IDRI | TCIA |
---|---|---|---|
Real | MSD Task 06 | – | 3.987 |
LIDC-IDRI | 3.987 | – | 4.744 |
TCIA | 1.858 | 4.744 | – |
Synthesis | HA-GAN | 98.208 | 116.260 |
MAISI | 19.008 | 31.370 | 20.338 |
表1:MAISI模型和基线方法使用其发布的检查点的Fréchet起始距离得分,以多个公共数据集为参考
随后,我们使用我们的数据集重新训练了几种最先进的基于扩散模型的方法。表2和表3中的结果表明,对于我们的数据集和未发现的数据集(autoPET 2023),我们的方法始终优于以前的方法。
Method | FID (XY Plane) ↓ | FID (YZ Plane) ↓ | FID (ZX Plane) ↓ | FID (Average) ↓ |
---|---|---|---|---|
DDPM | 10.031 | 36.782 | 43.109 | 29.974 |
LDM | 12.409 | 19.202 | 22.452 | 18.021 |
HA-GAN | 10.439 | 10.108 | 10.842 | 10.463 |
MAISI | 1.225 | 2.846 | 2.854 | 2.308 |
表2:以我们的数据集为参考,比较我们的基础模型和重新训练的基线方法之间的Fréchet起始距离得分
Method | FID (XY Plane) ↓ | FID (YZ Plane) ↓ | FID (ZX Plane) ↓ | FID (Average) ↓ |
---|---|---|---|---|
DDPM | 18.524 | 23.696 | 25.604 | 22.608 |
LDM | 16.853 | 10.191 | 10.093 | 12.379 |
HA-GAN | 17.432 | 10.266 | 13.572 | 13.757 |
MAISI | 14.165 | 5.770 | 8.510 | 9.481 |
表3:以autoPET 2023为参考,比较我们的基础模型与重新训练的基线方法之间的Fréchet Inception Distance得分
图3显示,我们的方法生成的图像显示出显著增强的细节和更准确的全局解剖结构。
图3:基线方法(使用我们的大规模数据集重新训练)和我们的方法之间生成图像的定性比较
生成模型最重要的应用之一是合成新数据,以用于模型训练中的数据扩充。我们可以通过评估包含合成数据的影响来评估生成图像的质量。我们采用了Auto3DSeg管道,这是一种在MONAI中开发医学图像分割解决方案的自动管道,并从头开始训练每个分割模型,以减少随机性通过五倍交叉验证。
有两组实验:
- 真实:正常的模型训练是在真实数据上进行的。
- 真实+合成:在训练过程中,真实数据和合成数据以相等的比例组合,以显示合成数据对数据扩充的效果。
如表4所示,所有合成数据在五种肿瘤类型中都对测试集的最终性能产生了积极影响(约2.5%~4.5%的改善)。这些结果表明,使用合成数据训练的模型具有更好的可推广性。
Experiment | Dataset | Tumor Type | Dice Score | Improvement |
---|---|---|---|---|
Real | MSD Task 06 | Lung Tumor | 0.581 | – |
Real + Synthetic | 0.625 | 4.5% | ||
Real | MSD Task 10 | Colon Tumor | 0.449 | – |
Real + Synthetic | 0.490 | 4.1% | ||
Real | In-House Bone Lesion | Bone Lesion | 0.504 | – |
Real + Synthetic | 0.534 | 3.0% | ||
Real | MSD Task 03 | Hepatic Tumor | 0.662 | – |
Real + Synthetic | 0.687 | 2.5% | ||
Real | MSD Task 07 | Pancreatic Tumor | 0.433 | – |
Real + Synthetic | 0.473 | 4.0% |
表4:与不同肿瘤类型的各种模型的基线性能相比,Auto3DSeg的平均骰子得分有所提高
定性评估
图4显示了三例异常病例的定性评估,可以看出MAISI在正常器官和异常肿瘤区域都能产生良好的CT生成质量,如每个子图的方框所示。我们的结果表明MAISI能以高保真度有效地描绘异常组织边界,证明了其在医学成像中基于分割掩模条件捕捉复杂细节方面的稳健性,MAISI有可能有效地增强生成的CT图像的多样性和真实性用于数据增强目的。
图4:结肠肿瘤(上)、骨病变(中)和肺肿瘤(下)的分割掩模(左)和生成的CT图像(右)的示例
值得注意的是,在每种情况下,MAISI都准确地模拟了异常肿瘤区域的出现,并为丰富肿瘤形态和空间分布变化的数据集开辟了可能性,这些发现突出了MAISI作为增强医学成像数据集的强大工具的潜力,从而提高了机器学习模型在临床应用中的稳健性和通用性。
总结
MAISI是一个最先进的基础人工智能模型,用于生成具有相应标签的3D高分辨率合成医学图像,以解决数据限制、降低注释成本和维护患者隐私。凭借其实现高质量分辨率和分割127个解剖类别的能力,MAISI有望在医学成像领域产生重大影响,将MAISI生成的合成数据纳入训练分割模型已经证明了显著的性能改进,为临床应用中提高鲁棒性和泛化能力铺平了道路。
探索使用MAISI为您的项目生成合成数据的潜力,请加入早期访问计划。
鸣谢
所有合著者都希望指出,他们对本文的研究和本文的写作做出了同等的贡献。