Command Palette
Search for a command to run...
DiffNR:用于稀疏视图 3D 断层扫描重建的扩散增强神经表示优化
DiffNR:用于稀疏视图 3D 断层扫描重建的扩散增强神经表示优化
Shiyan Su Ruyi Zha Danli Shi Hongdong Li Xuelian Cheng
摘要
神经表示(Neural Representations, NRs),例如神经场(neural fields)和 3D 高斯(3D Gaussians),能够有效地对计算机断层扫描(CT)中的体数据进行建模,但在稀疏视图(sparse-view)设置下会面临严重的伪影问题。为了解决这一问题,我们提出了 DiffNR,这是一个利用扩散先验(diffusion priors)来增强 NR 优化的创新框架。该框架的核心是 SliceFixer,这是一个旨在修复退化切片中伪影的单步扩散模型。我们在网络中集成了专门的条件层(conditioning layers),并开发了定制的数据策展(data curation)策略以支持模型微调。在重建过程中,SliceFixer 会定期生成伪参考体数据(pseudo-reference volumes),通过提供辅助的 3D 感知监督来修复约束不足的区域。与以往将 CT 求解器嵌入到耗时的迭代去噪过程中的方法相比,我们的“修复与增强”(repair-and-augment)策略避免了频繁调用扩散模型,从而实现了更好的运行性能。广泛的实验表明,DiffNR 的峰值信噪比(PSNR)平均提升了 3.99 dB,在不同领域间具有良好的泛化能力,并保持了高效的优化性能。
一句话总结
为了解决稀疏视图 3D 断层扫描重建中的伪影问题,所提出的 DiffNR 框架通过集成 SliceFixer 增强了神经表示优化。SliceFixer 是一个单步 diffusion 模型,通过一种高效的“修复与增强”策略定期生成伪参考体数据以进行 3D 感知监督,使平均 PSNR 提升了 3.99 dB。
核心贡献
- 本文引入了 DiffNR,这是一种用于稀疏视图 3D 断层扫描重建的新型优化框架,通过条件 diffusion 先验来增强神经表示。
- 本研究提出了 SliceFixer,这是一种专门设计的单步 diffusion 模型,通过专门的条件层和定制的数据策划策略来纠正重建 CT 切片的伪影。
- 所提出的“修复与增强”策略利用 SliceFixer 定期生成伪参考体进行 3D 感知监督,在保持高优化效率的同时,实现了 3.99 dB 的平均 PSNR 提升。
引言
稀疏视图 CT 重建对于减少医学成像和安全检查过程中的辐射暴露至关重要。虽然神经场和 3D Gaussians 等神经表示 (NR) 方法能够提供高质量的体数据建模,但在投影数据有限时,它们往往会在约束不足的区域产生严重的伪影。现有的使用 diffusion 模型解决该问题的神经先验方法通常面临处理时间长、切片间抖动以及幻觉问题。本文通过名为 DiffNR 的新型框架利用“修复与增强”策略。引入了 SliceFixer,这是一种经过微调的单步 diffusion 模型,用于纠正重建切片中的伪影,并定期生成伪参考体以提供 3D 感知监督。该方法避免了频繁进行 diffusion 查询带来的计算负担,同时显著提高了重建精度和体数据一致性。
数据集
通过合成由含有大量伪影的重建结果与干净的地面真值 (ground truth) 体数据组成的配对切片,开发了一个专门用于训练 SliceFixer 的数据集。数据集的构成和处理包括以下内容:
- 数据来源与合成:利用公开的 3D CT 体数据作为基础。使用断层扫描工具箱,在完整的 360∘ 角度范围内合成 K 个密集投影,以创建高质量的地面真值数据。
- 伪影模拟:为了创建配对的伪影样本,通过随机采样密集投影的子集来模拟稀疏视图场景。采用均匀和非均匀的视图分布来引入多种伪影模式,从而提高模型的鲁棒性。
- 欠拟合策略:为了增加训练样本的难度和多样性,刻意对神经重建模型进行欠拟合。通过将优化限制在标准训练迭代次数的 25% 到 50%,产生因收敛不完全而导致明显伪影的重建结果。
- 数据集混合:最终的训练集采用混合神经表示方法。将神经场和 3D Gaussians 的重建结果按 1:1 的比例结合,以防止 diffusion 模型过拟合于单一的重建模式,并鼓励学习泛化先验。
方法
提出 SliceFixer,这是一种基于 diffusion 的模型,旨在细化由神经表示 (NRs) 产生的损坏 CT 切片,并将其集成到更广泛的优化框架中以增强体数据重建。整体方法利用了 diffusion 模型进行高保真图像生成的优势,同时解决了独立后处理在维持体数据一致性方面的局限性。
SliceFixer 基于单步 diffusion 模型 SD-Turbo 构建,以实现高效推理。如框架图所示,其架构由变分自编码器 (VAE) 和 U-Net 组成。VAE 将来自 NR 的损坏切片编码为潜在表示 (latent representation)。该潜在表示连同条件信息和去噪时间步一起输入 U-Net,由 U-Net 预测目标潜在特征。随后,通过使用 VAE 解码器对预测的潜在表示进行解码,从而重建出细化后的切片。
该模型以双平面 X 射线投影和文本提示 (text prompt) 为条件,以引导细化过程。双平面投影提供全局结构线索,而文本提示提供高层语义指导。投影的图像特征使用预训练的 RAD-DINO 编码器提取,并通过交叉注意力机制与文本嵌入聚合,形成 diffusion 模型的最终条件输入。
为了适配预训练的 SD-Turbo 模型,采用了微调策略,将可训练的 LoRA 适配器注入 VAE 和 U-Net 模块,并为跳跃连接引入零卷积层。这使得模型能够在保留基础模型丰富视觉先验的同时,学习 CT 切片修复这一特定任务。训练目标集成了多种损失函数:L2 损失、LPIPS 损失、CLIP 对齐损失、基于 GAN 的对抗损失以及结构相似性 (SSIM) 损失,以确保感知质量和解剖保真度。
进一步将该框架扩展为名为 DiffNR 的 diffusion 增强型 NR 优化流水线。该方法并非将 SliceFixer 作为后处理步骤,而是将其作为 NR 训练循环中的主动组件。如流水线图所示,首先使用标准图像损失和低级正则化来优化神经表示。定期从当前模型中查询一个体数据,并将其切片通过 SliceFixer 处理以生成伪参考体。该参考体随后以 3D SSIM 损失的形式提供额外的 3D 监督,有助于维持体数据一致性并减少切片间抖动。这种“修复与增强”策略允许模型利用 SliceFixer 生成的高质量、无伪影切片作为一种强大的增强监督形式,在保持训练效率的同时提高整体重建质量。
实验
实验使用 ToothFairy 和 LUNA16 数据集对所提出的 DiffNR 方法进行评估,并结合分布外 (out-of-distribution) 测试和下游肺部分割任务来验证重建质量和泛化能力。结果表明,DiffNR 有效抑制了传统方法和基于 diffusion 方法中常见的伪影和幻觉,同时恢复了更精细的结构细节。此外,该方法通过生成在医学分割任务中产生更准确结果的重建体数据,展示了强大的实际应用价值。
在分布外数据集上评估 DiffNR 方法,并将其与几种先进的重建方法进行比较。结果显示,DiffNR 在所有视图设置下均实现了最高的性能,显著优于基准 R2-Gaussian 方法,并超越了其他基于 diffusion 和迭代的方法。该方法有效地抑制了伪影和幻觉,展示了强大的泛化能力和重建质量。DiffNR 在所有视图设置下的分布外数据集上均取得了最佳性能。DiffNR 显著改进了 R2-Gaussian 基准,表现出强大的泛化能力和伪影抑制能力。DiffNR 优于其他基于 diffusion 和迭代的方法,表明其具有卓越的重建质量和鲁棒性。
在多个数据集和视图设置下,将 DiffNR 与几种现有的重建方法进行比较。结果表明,DiffNR 在重建质量和伪影抑制方面表现优异,特别是在分布内和分布外场景中,关键指标均优于基准方法。该方法在分割等下游应用中展示了强大的泛化能力和有效性,在与地面真值的一致性方面优于其他方法。DiffNR 在不同的视图设置和数据集上均优于现有方法,在 Dice 和 ASD 指标上均获得了更高的分数。DiffNR 展示了强大的泛化能力,在其他方法难以处理伪影和幻觉的分布外数据上改善了结果。该方法与地面真值体数据保持高度一致,与基准方法相比,带来了更准确的下游分割结果。
进行了消融实验,以评估 DiffNR 方法中不同组件对重建质量的影响。结果表明,将 SliceFixer 与 SSIM 损失相结合并将其作为优化流水线的一部分,可以获得最佳性能,与基准方法相比,显著提高了 PSNR 和 SSIM 指标。将 SliceFixer 与 SSIM 损失集成比单独使用它能提高重建质量。将 SliceFixer 作为独立的后处理步骤会导致切片抖动和幻觉。在优化流水线中将 SliceFixer 与 SSIM 损失结合,在 PSNR 和 SSIM 上均达到了最高性能。
进行了消融实验,以评估 SliceFixer 和 DiffNR 中不同设计选择对重建质量的影响。结果表明,在微调期间使用更高分辨率并结合 SSIM 损失可以提高性能,而添加双平面投影可以进一步增强重建。将 SliceFixer 集成到优化流水线中至关重要,且与 3D 感知损失相比,体素级 L1 损失会导致性能下降。在更高分辨率下并结合 SSIM 损失微调 SliceFixer 可以提高重建质量。添加双平面投影作为条件输入,通过提供更丰富的结构线索来提升性能。将 SliceFixer 集成到优化流水线中可以防止在独立后处理中观察到的切片抖动和幻觉。
在两种数据集的不同视图条件下,将 DiffNR 与各种基准方法进行了对比评估。结果表明,DiffNR 在重建质量方面优于其他方法,特别是在 PSNR 和 SSIM 指标方面,同时与之前的基于 diffusion 的方法相比,保持了合理的计算效率。该方法较基准神经表示方法实现了持续改进,并在显著缩短处理时间的同时达到了具有竞争力的性能。DiffNR 实现了比传统方法和基于 diffusion 方法更高的重建质量,特别是在 PSNR 和 SSIM 指标上。DiffNR 明显快于之前的基于 diffusion 的方法,将处理时间从数小时缩短至数分钟。DiffNR 持续改进了基准神经表示方法,在不同数据集和视图条件下表现出鲁棒的性能。
在多个数据集和视图设置下,将 DiffNR 与各种先进的重建方法进行对比,以测试泛化能力和下游分割精度。消融实验进一步验证了将 SliceFixer 与 SSIM 损失及双平面投影集成到优化流水线中以防止伪影和幻觉的重要性。总体而言,与现有的基于 diffusion 的方法相比,DiffNR 展示了卓越的重建质量、对分布外数据的强健泛化能力以及显著提高的计算效率。