原文链接:Anatomically constrained and attention-guided deep feature fusion for joint segmentation and deformable medical image registration
一、Hilghtlights:
- 利用分割可以在配准过程中培养先验知识。
- 多任务学习策略提高了配准性能。
- 所提出的解剖约束策略对预测进行了细化。
- 跨任务注意力的特征传递最大限度地提高了信息的流动性。
- 注意力机制可以将注意力重新集中到变形区域。
思考:利用分割是从voxlemorph就有的技巧,我和导师讨论的时候老师说可以考虑构建一个一体化的架构,端到端的完成分割和配准,现在看来已经有人完成了啊┑( ̄Д  ̄)┍;tranformer开始进入配准领域了,实际上注意力机制是很适配配准问题的,我的理解是让模型重点关注大变形、复杂变形区域,这样可以提高性能
二、摘要
可变形图像配准的解剖学合理结果的主要目的是通过最小化一对固定图像和运动图像之间的差异来提高模型的配准精度。由于许多解剖学特征彼此密切相关,因此利用辅助任务(例如有监督解剖分割)的监督标签有可能在配准时增强扭曲图像的真实感。在这项工作中,我们采用多任务学习框架将配准和分割表述为一个联合问题,我们利用辅助监督分割的解剖约束来增强预测图像的真实性。首先,我们提出了一个跨任务注意力块来融合来自配准和分割网络的高级特征。在初始解剖分割的帮助下,配准网络可以从学习任务共享特征相关性并快速关注需要变形的部分中受益。另一方面,将初始翘曲图像的GT标注和预测出的分割图的解剖分割差异整合到损失函数中,以指导配准网络的收敛。理想情况下,良好的变形场应该能够最小化配准和分割的损失函数。从分割中推断出的体素解剖约束有助于配准网络达到可变形和分割学习的全局最优值。在测试阶段,这两个网络都可以独立使用,当分割标签不可用时,只能预测配准输出。
重点:额外的监督解剖分割可以帮助我们明确学习任务之间的相关性,以提高最终的可变形配准精度,学习相关性的动机是,两个任务之间的相关性在解剖学先验中比单独来看更优。正如之前的工作中所讨论的,医学图像配准和分割是互补和相互关联的(即,被图像配准算法扭曲的移动注释标签经常用于图像分割)。在存在高解剖变异性的情况下,配准和分割之间的这些任务鲁棒相关性具有在很大程度上提高目标医学图像配准精度的潜力。
在本文中,我们重点处理了JRS()框架中有监督分割网络和无监督配准之间的信息流。我们利用这两项任务之间的解剖学联系来提高配准的准确性。为了将全局约束形式的解剖学先验纳入可变形学习过程,我们建议可以通过两种方式利用这种相关性。一方面,我们提出了明确的方法,学习配准和分割之间的任务特征相关性。这是通过使用我们新设计的跨任务注意力块来模拟配准和分割特征之间的相互作用和互补性来实现的。目标配准任务可以利用从分割中获得的全局解剖上下文信息流,在翘曲后生成更真实和解剖学上合理的图像。另一方面,这种相关性被用于增强配准网络产生的变形场,并在变形的运动图像和固定图像之间加强解剖学对应关系。我们通过计算真实固定注释标签与初始配准后分割网络从扭曲图像预测的标签之间的差异来近似解剖约束。由于配准和分割是耦合的,我们假设估计的解剖学约束可以从分割转移到配准。我们建议利用这种关系来指导目标配准任务的变形场细化。结合这两种相关性利用方式,我们提出了解剖学相关性感知可变形医学图像配准(AC-DMiR)方法。
三、方法
对于第一种策略,局部体素水平损失函数(Ferrante等人,2018,Ferrante等人,2017)没有纳入全局信息,因此不一定生成更强的解剖学合理性。我们的方法采取后一种策略。我们利用监督分割和任务相关性来改善两个任务之间的信息流,并在变形后产生更逼真的图像。
在这项工作中,我们建议使用监督分割网络来学习解剖学先验知识。使用监督分割允许我们利用两个任务之间的相关性来进一步改善扭曲运动图像和固定图像之间的解剖映射。首先,我们设计了跨任务注意力块来学习配准和分割之间的任务间特征相关性,并将其转移到目标配准任务中。在我们的实现中,我们利用了一个联合配准和分割框架,该框架在学习过程中通过联合损失进行了优化。此外,通过计算分割网络的初始预测与固定标签图之间的差异,提出了一种解剖约束策略。由于配准和分割是自然耦合的,我们不仅在分割中使用解剖学约束,而且在配准网络中也利用解剖约束作为损失函数。这鼓励网络参数收敛到最终扭曲运动图像具有高解剖学合理性的点
3.1 Anatomical correlation-aware architecture
我们调整了我们的相关性感知 AC-DMiR 框架,以利用有关解剖结构之间空间关系的信息,利用 MTL(多任务学习) 框架的最新进展进行配准和分割任务(Xu 等人,2018 年,Vandenhende 等人,2020 年)。使用两个混合变压器-CNN 骨干网络(Chen 等人,2022 年)对特定于任务的初始配准和分割进行建模,以生成初始预测。下一步是使用任务共享解剖相关模块来明确学习配准和分割之间的相关性,并合并来自其他任务的补充信息以改进最终配准预测。
3.2 Task-specific initial predictions
3.2.1 无监督配准
首先生成初始预测,以便了解配准和分割信息之间的相关性。此时,我们的配准网络旨在计算初始变形场,将和对齐。
相似性度量 这里图像是同模态的,强度分布相似。因此可以使用MSE和LNCC来衡量(比较常见喽):
正则化 为了保证平滑,引入扩散正则器、弯曲能量、抗折叠损失(也比较常见喽,如果对配准模型的损失函数有不理解的同学,请看我的博客:学习笔记:医学图像配准简介—附voxelmorph模型,或者阅读原文,有完整的解释)
辅助信息损失 良好的变形场可以高精度地预测解剖对应关系,这意味着变形后的source和target应该很好地重叠。使用DICE分数计算分割掩码的重叠情况:
当然,直接使用这个损失方程式不足以实现自动区分网络(auto-differentiation networks)的。原文中使用了一种近似方法。感兴趣的可以阅读一下。
3.2.2 有监督分割
鉴于靶器官变形与周围器官特性基本相关,我们考虑利用解剖学信息来限制最终的配准过程。分割网络的输入时初始配准网络的输出,也就是,该网络预测对应输出的标签掩码
这里有点复杂,贴原文:
我的理解是:有K个器官,对每个器官,target图都会有GT标注,现在对初始配准后的source图做分割(也就是要对 的每一个体素做一个onehot编码),GT就是target的GT,这样相当于把目标图像的分割信息融入了。
3.2.3 任务共享的解剖相关性模块(本文重点
为了学习配准和分割之间的相关性,需要将初始阶段的配准和分割出的高级特征输入到任务共享解剖相关模块(task-shared anatomical correlation module)中。这是通过跨任务注意力块(CTAB)来实现的,该块捕获了相互关系,并结合了来自初始配准和分割网络的高级表示的瓶颈特征(bottleneck features):This is accomplished through the implementation of the proposed Cross-Task Attention Block (CTAB), which captures the mutual relationship and incorporates bottleneck features of high-level representations from both initial registration and segmentation networks. 解剖学相关模块的设计受到多任务学习领域研究的启发(Vandenhende 等人,2020 年),其中使用了可比较的注意力模块来促进各种任务的联合学习。与Wang等人(2018)中提出的非局部注意力模块Non-Local Neural Networks不同,该模块是在单个图像上计算自我注意力,而本文所提出的CTAB试图模拟来自执行不同任务的两个网络的高级特征之间的解剖对应关系,也就是需要从两个图像的特征上计算注意力。
Tip:本文的公式乍一看不知道在干嘛,我的理解是通过三个不同的线性嵌入层获取向量——也就是两个位置的高级特征,然后用exp函数计算解剖相关性(在每个位置j上)。最后的输出f是需要归一化的。最后,由所有f组成的矩阵就包含了从A到每一个T中位置的非局部信息。公式9的思路来源于Non-local Neural Networks | IEEE Conference Publication | IEEE Xplore:
问题是为什么exp函数可以计算解剖相关性? ??还没搞懂
NEWBING yyds!
整个模块的输出是三种解剖相关性(特征):
通过这样的设计,我们能够促进配准和分割任务之间的高级信息流。一方面,CTAB以全局约束的形式整合分割任务中的解剖学先验,同时排除冗余信息。因此,我们明确地学习了这两个任务之间的相关性。另一方面,我们的目标是通过将注意力从分割网络共享到配准网络来获得更强大和可推广的解剖相关性。
TODO:这部分内容还没太看懂,需要加深理解。
3.2.4 专注于配准的最终解码器
通过上面的模块获取了上面的三种特征后,使用第三个(融合)特征,可以进行最终预测,与开始的初始预测相比,只是改变了输入——这次的输入是融合特征。
整个网络的损失为:
注:什么是解剖约束?给定一对移动和固定的图像输入。首先,比较初始分割解码器预测的分割图和真实固定标注的体素预测差异;然后将差异用作体素解剖约束指标。由于这两个任务是相互关联的,我们假设解剖学先验可以从分割转移到配准任务。具有高分割图预测差异的体素表明该特定点的解剖相关性较弱,因此应分配更高的权重至目标配准任务中。
3.2.5. Deformation field refinement with anatomical constraint,解剖约束策略
这一部分是介绍如何利用分割网络预测的分割图和固定图的标注之间的差异来约束变形场的优化过程,从而提高变形后图像的解剖真实性。也就是所说的解剖约束策略,展开来说,就是指作者利用分割网络的预测结果来指导配准网络的收敛。具体来说,作者计算了初始变形后的移动图像的分割预测和固定图像的分割真值之间的差异,作为一种体素级别的解剖约束,然后将这个约束加入到配准网络的损失函数中,以期望变形场能够最小化这个约束。这样做的目的是让变形后的移动图像和固定图像在解剖结构上有更好的对应关系,从而提高配准的精度和真实性。
重点是如何计算体素权重,也就是在计算损失函数时,对不同的体素赋予不同的权重,以反映它们的重要性或难度。如果一个体素在分割预测和真值之间有很大的差异,说明这个体素的解剖关联性较弱,也就是说它没有很好地对齐到正确的结构上。因此,作者给这个体素分配了一个较高的权重,让它在配准损失函数中有更大的影响,从而促使变形场对这个体素进行更多的调整,以达到更好的配准效果。
4、实验效果
做了六组消融实验,baseline是transmorph,是使用具有无限感受野的混合transformer-CNN模型,可以更准确地理解运动图像和固定图像体素之间的远程空间关系。感兴趣的同学可以阅读论文,或者看我的博客了解:
医学图像配准论文学习——TransMorph,用于无监督医学图像配准的transform_仿射网络_研究僧_GUO的博客-CSDN博客1
具体结果略。