分享自:

扩散概率模型用于三维医学图像生成

期刊:Scientific ReportsDOI:10.1038/s41598-023-34341-2

本研究的通讯作者是Daniel Truhn博士,其单位是德国亚琛大学医院诊断与介入放射科。主要作者包括Firas Khader、Gustav Müller-Franzes等多位研究人员,他们分别来自德国亚琛大学医院、RWTH亚琛大学、Ocumeda GmbH、海德堡大学医院、维尔茨堡大学医院、美因茨大学等机构。这项研究成果以《Denoising Diffusion Probabilistic Models for 3D Medical Image Generation》为题,于2023年发表在科学期刊《Scientific Reports》上。

在学术背景方面,本研究属于医学影像人工智能(AI)与计算机视觉交叉领域,具体聚焦于生成式模型在三维(3D)医学图像合成中的应用。近年来,扩散概率模型(Diffusion Probabilistic Models)在自然图像生成领域取得了突破性进展,例如DALL-E 2和Stable Diffusion等项目已能根据文本输入生成逼真的图像。然而,医学影像数据通常以三维体积的形式存在(如磁共振成像MRI和计算机断层扫描CT),扩散模型在此类数据上的应用和系统性评估此前尚未充分开展。合成医学图像在隐私保护型人工智能(即在无需共享原始患者数据的情况下进行协作研究)、以及在小数据集上进行数据增强以提高模型性能方面,具有至关重要的潜在价值。生成对抗网络(GANs)曾是生成3D医学图像的主流方法,但其存在训练复杂、易发生模式坍塌(生成的样本缺乏多样性)、以及难以与文本等非图像模态关联等局限性。相比之下,扩散模型在非医学领域已展现出生成高质量、多样化图像并链接不同数据模态的强大能力。因此,本研究旨在系统性地探索扩散模型在医学领域生成三维数据的潜力,目标是开发一种能够生成高质量、解剖结构正确且多样的3D MRI和CT图像的架构,并验证这些合成数据在临床实际应用(如在数据稀缺场景下提升分割模型性能)中的价值。

研究的详细工作流程包含几个核心步骤:数据准备与预处理、构建并训练矢量量化生成对抗网络(VQ-GAN)以学习图像的压缩潜在空间表示、在潜在空间上训练3D扩散概率模型进行图像生成、以及利用合成数据进行下游任务(如分割)的应用验证。

首先,在数据准备阶段,研究团队选择了四个公开可用的医学影像数据集,涵盖了不同的解剖部位和成像模态,以证明方法的普适性。这些数据集包括:阿尔茨海默病神经影像倡议(ADNI)的脑部MRI数据(998个研究)、癌症影像存档库(TCIA)中肺癌影像数据库联盟(LIDC)的胸部CT数据(1010个研究)、杜克大学(Duke)的乳腺MRI数据(1844个研究)以及斯坦福大学MRNet的膝关节MRI数据(1250个研究)。这些数据集规模相对较小(均≤1844个样本),旨在专门探索模型在“小数据”场景下的表现。预处理过程针对每个数据集进行了标准化,包括图像重采样至标准体素间距、强度归一化到[-1, 1]范围、中心裁剪或调整为统一维度(如64×64×64, 128×128×128, 256×256×32)等操作,并采用了垂直翻转进行数据增强。

其次,研究采用了一种新颖的两阶段架构来应对3D医学图像数据量大、计算负担重的挑战。第一阶段,训练一个3D VQ-GAN自编码器。这是一个关键的技术创新点。VQ-GAN通过编码器将高分辨率3D图像压缩到一个低维度的潜在空间,然后使用一个学习到的“码本”对潜在特征向量进行量化,最后通过解码器重建图像。与传统的变分自编码器(VAE)相比,VQ-GAN结合了对抗性训练(使用了一个结合切片判别器和3D判别器的判别器),能生成更清晰、细节更丰富的重建图像,避免了模糊问题。研究团队将原有2D VQ-GAN架构扩展到了3D,用3D卷积替换了2D卷积。他们发现潜在空间的压缩因子至关重要:实验表明,当压缩因子为4(例如,将256×256×32的图像压缩到64×64×8的潜在空间)时,能很好地保留解剖细节;而压缩因子为8会导致关键解剖特征丢失。训练好的VQ-GAN的码本和编码器-解码器将用于后续阶段。

第三阶段,在VQ-GAN学习到的量化潜在空间(经归一化处理后)上,训练一个3D扩散概率模型。这是另一个核心创新。扩散模型通过一个前向过程逐步向数据添加噪声,然后训练一个U-Net网络学习逆向去噪过程,从而可以从纯噪声中生成新的数据样本。为了处理3D数据,研究团队设计了一个专门的3D U-Net架构。该架构在编码器和解码器的每个模块中,依次使用3D卷积层、空间注意力层(Spatial Attention Layer)和深度注意力层(Depth-wise Attention Layer)。空间注意力层使高分辨率图像平面上的每个元素能够关注平面上的其他元素,而深度注意力层则允许同一平面位置的特征关注不同深度切片上的特征。这种设计使得模型能够有效捕捉3D体积内跨切片的连续性和一致性信息。训练完成后,通过从标准高斯分布采样噪声,并运行训练好的扩散模型进行逆向去噪,可以在潜在空间中生成新的隐式表示,再通过VQ-GAN的解码器将其解码为最终的3D合成图像。

第四阶段,应用验证。为了展示合成数据的实用价值,研究设计了一个实验:使用在杜克大学乳腺MRI数据集上训练的扩散模型,生成了2000张合成乳腺MRI图像(无任何真实标签)。然后,利用这些合成图像,以自监督学习的方式对一个Swin UNETR分割模型进行预训练,预训练任务被设置为图像修复(Inpainting)。之后,在另一个内部数据集(包含200个带有手工标注乳腺轮廓分割掩模的真实乳腺MRI图像)上,使用不同比例(5%, 10%, 20%, 40%, 80%, 100%)的真实数据对预训练模型进行微调,执行乳腺分割任务。作为对照,另一个Swin UNETR模型不经过合成数据预训练,直接在不同比例的真实数据上训练。通过比较两组模型在测试集上的分割性能(采用Dice相似系数衡量),来评估合成数据预训练的效果。

研究取得了一系列重要结果。在图像生成质量方面,合成图像在视觉上呈现出高度的真实感。更重要的是,为了超越传统自动化指标评估的局限性,研究邀请了两名分别拥有9年(读者A)和5年(读者B)经验的放射科医生,对每个数据集生成的50张合成图像进行人工盲法评分,评分维度包括“真实的图像外观”、“切片间一致性”和“解剖正确性”。评分结果非常积极:读者A认为189/200张图像至少“基本真实,仅存在次要不真实区域”;191/200张图像在“大多数切片(>10片)一致”;185/200张图像显示“仅存在次要或没有解剖不一致性”。读者B的评分结果类似。这从临床专家视角强有力地证明了本方法生成的3D合成图像具有高度的逼真度和解剖合理性。

在与传统方法的对比结果中,研究将本扩散模型与一个已建立的Wasserstein GAN with Gradient Penalty (WGAN-GP)模型在脑部MRI数据上进行了比较。使用多尺度结构相似性指数(MS-SSIM)来衡量生成样本的多样性(数值越低表示多样性越高)。结果显示,GAN模型的MS-SSIM分数高达0.9996,表明其生成的图像几乎完全相同,多样性极低,存在严重的模式坍塌问题。而扩散模型的MS-SSIM分数为0.8557,更接近原始真实数据的分数(0.8095),证明扩散模型能够生成更多样化、更能代表原始数据分布的样本。

在下游应用验证结果中,实验数据清晰地表明,在数据稀缺场景下,使用来自其他机构的合成数据进行预训练能显著提升模型性能。特别是在仅使用5%内部真实数据(即约10个样本)进行微调时,经过合成数据预训练的模型其乳腺分割的Dice分数达到了0.95,而未经过预训练的模型Dice分数仅为0.91。随着可用真实数据比例的增加,预训练带来的优势逐渐减小,但在低数据区域(5%-40%)改善始终明显。这证明了合成数据可以有效作为知识迁移的载体,帮助模型在目标数据有限的情况下学习到更鲁棒的视觉特征。

基于以上结果,本研究得出的结论是:基于潜在扩散模型的方法能够成功生成高质量、解剖结构正确且多样的三维医学图像(MRI和CT),其性能优于传统的生成对抗网络(GANs),特别是在避免模式坍塌和生成多样性方面。该方法即使在相对较小的数据集(约1000个样本)上也能稳健训练。更重要的是,这些合成数据具有切实的应用价值,能够通过自监督预训练的方式,有效提升下游医学图像分析模型(如分割模型)在数据稀缺情况下的性能,为解决医学影像AI中常见的小样本学习问题提供了新的可行路径。

本研究的亮点在于:第一,这是首次对潜在扩散模型在多种三维医学影像模态(MRI和CT)及不同解剖部位上进行的大规模系统性评估和验证。第二,在方法学上具有显著的创新性,包括将VQ-GAN成功扩展到3D医学图像,并设计了结合空间与深度注意力的专用3D U-Net扩散模型架构,有效建模了3D体积数据的空间连续性。第三,评估方式全面且令人信服,不仅采用了定量指标,更重要的是引入了放射科医生的人工盲法评分作为图像质量的“金标准”,极大地增强了结果的可信度与临床相关性。第四,超越了单纯的生成任务,通过一个精心设计的分割实验,清晰、量化地证明了合成数据在真实临床AI应用场景(跨机构协作、小数据学习)中的实用价值,将研究工作从技术探索延伸到了价值验证。

此外,研究还公开了全部代码,以促进该领域的进一步研究并作为未来工作的基准。同时,作者也讨论了本工作的局限性,例如使用的公开数据集分辨率并非当前临床诊断的最高标准,以及模型在更大规模、更高分辨率数据上的潜力有待未来在算力更充足时进行探索。这些都为后续研究指明了方向。该项研究为三维医学图像合成领域树立了一个新的标杆,展示了扩散模型在该领域的巨大潜力和优越性,为医学影像的隐私保护、数据增强和AI模型开发提供了强大的新工具。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com