https://github.com/multimodallearning/pdd_net

Closing the Gap between Deep and Conventional Image Registration using Probabilistic Dense Displacement Networks

摘要

诊断任务，手术图像引导，放射治疗以及运动分析在很大程度上依赖于准确的患者内对齐。患者间的配准可以依赖基于atlas风格或者landmark定位或者形状分析。当标签少并且解剖差异较大，传统的注册方法往往仍然优于深度学习方法，到目前为止，深度学习方法主要处理相对较小或较低复杂性的变形。我们通过利用概率密集置换优化的思想来解决这个缺点，这种思想在许多大变形的配准任务中表现出色。我们建议设计一个具有近似最小卷积和平均场推理的网络，用于在离散弱监督配准设置下的微分位移正则化。通过使用这些有意义并理论上已证明的约束，我们可学习的配准算法包含很少的可训练权值(主要用于特征提取)，并且通过很少的标记扫描更容易训练。它在训练和推断方面非常快，并且在腹部CT具有挑战性的患者间配准方面达到了最先进的准确性，比以前的深度学习方法有15%的Dice重叠。

关键字：registration、deep learning、probabilistic、 abdominal

1、简介和相关工作

传统的医学图像配准主要依赖于迭代，最小化不同多尺度的翘曲waping moving相似度的和正则化惩罚来实现。基于深度学习的图像配准(Deep learning based image registration, DLIR)旨在模拟这一过程，通过训练一个卷积网络，该网络可以预测给定两次新的未训练扫描的非线性对齐函数。因此，一个单一的前馈传递函数必须使用许多卷积层而不是多个warping步骤。DLIR的监督可以基于自动或手动对应、语义标签或内在成本函数。它在时间敏感的应用方面具有巨大的潜力，例如通过多atlas配准的图像引导、融合、跟踪和形状分析。然而，由于大空间的潜在变形，可以映射两个对应的解剖到另一个，这个问题比图像分割的约束少得多，因此仍然是一个开放的挑战。

脑部局部变形为几毫米，并有大量的标记数据集。对于腹部、前列腺或肺部的其他解剖，形状变化只有几厘米，DLIR主要应用于不太复杂的病例。对于吸气-呼气肺配准，DLIR的准确性~2.5mm仍低于常规方法<1mm。弱监督的DLIR方法Label-Reg在腹部CT上，病患建对齐平均的Dice为42.7%比NiftyReg算法Dice56.1%要低。

我们的假设是，如果不借助复杂的多阶段翘曲管道warping pipelines，很难用深度连续回归网络对不同患者之间大的形变进行建模。相反，离散配准的使用，即同时探索一个大的量化位移空间，已经被证明可以更有效地捕捉腹部和胸部的变形[5,12,16]，并且可以通过很少或单个的扭曲步骤实现。不出所料，DLIR的2D视觉中已经探索了离散位移设置:即FlowNet-C[2]。提出了一种不包含可训练权值的相关层(见[2]中的Eq. 1)，通过使用产生441通道联合特征图的密集量化位移空间(21×21像素偏移量)移动移动图像，计算两幅图像的特征相似度。接下来，我们学习了一个非常大的441(+32)×256×3×3的核(然后进一步卷积)，它忽略了位移空间的显式四维几何。因此，大量的优化参数导致了对监督训练数据的巨大需求。将这个想法扩展到3D是非常困难的，因为在稠密相关之后，维数增加到6D，尽管它有很多好处，但是还没有被考虑到。概率和不确定性建模已经在DLIR中进行了研究，参考文献[9,17]，但没有在离散的环境中进行研究。

我们提出了一种新的DLIR学习模型，该模型通过引入带有可微约束的强正则化来更好地利用概率密集位移采样probabilistic dense dispalcement sampling 的优势问题的本质。因此，我们使用均值场mean-field推断进行正则化[8,18]，并使用近似最小卷积[3]计算标签间的兼容性，从而从空间变换的拟合中分离卷积特征。我们的特征提取器使用三维可变形卷积3D defoemable convolutions[4]，是非常轻量级的。据我们所知，这是将离散DLIR与均值场正则化的可微应用相结合的方法首选。与之前的工作相比，我们的模型需要更少的可训练重量，捕获更大的变形，并且可以从少量的标记扫描训练到高精度。我们同时提出了一种新的非局部标签损失来代替更广泛使用的基于空间变换spatial transformer的损失。

2、方法

我们旨在通过在和上约束并学习特征映射，找到空间变换，对齐和。为了训练处一个合适的特征提取，该特征提取对噪声和信息强度不敏感。在训练期间我们提供了一个监督标记，.我们把空间坐标定义为连续变量，使用三线性插值对离散网格进行采样。在粗糙的网格上用(几千个)控制点来参数化。位移范围被限制在一个离散的位移空间discrete displacement space，其线性间距为，其中q是定义捕获范围的标量，在我们的例子中，是3375。网络模型预测出位移概率displacement probabilities的6D tensor,每个控制点的维数4-6的和是1。概率与位移的内积得到这些概率估计数的加权平均值，从而得到在推理过程中的三维位移。

1）卷积特征学习网络:为了学习一个有意义的非线性映射，从输入强度到稠密的特征量(通道，步长为3)，我们采用Obelisk
方法,一个三维可变形卷积与可训练的偏移量，然后一个简单的1×1MLP非常有效地捕捉空间上下文。我们通过在Obelisk层之前添加一个具有4个通道的普通5×5×5卷积核来扩展作者的实现，从而学习边缘特征。该网络有64个空间滤波器偏移量，共有120k个可训练参数，用于固定和移动扫描，生成和。

2）致密位移不相关的相关层：第一部分提供了特征表示，我们旨在找到正则化为一场displacement field，将矢量vector 分配给每个控制点来进行非线性变换最大限度的提高了fixed和warped moving scan的（label）相似性。正如传统的离散配准，和相关层correlation layer中描述的那样，我们对位移搜索空间上的相似性度量进行密集的评估。通过学习描述符特征维度的负均方差（mean squared error, MSE）获得不相似dissimilarities 6D tensor 。可以使用差异度量例如相关系数correlation coefficient。由于控制点稀疏性，实验中对位移相似度displacement similarity的评价要求不差过2GFlops。位移捕捉范围q设置为0.4.

3）使用最小卷积和均值场推理正则化：由于非线性配准通常是不适定的ill-posed，因此采用附加先验保持形变空间平滑。对比其他在DLIR上的工作，原则上学习一个一个无约束的形变，并仅强制空间平滑作为损失项，我们建议将正则化约束建模为网络的一部分。基于位移平方差常用于Markov 随机场配准(Markov random field, MRF)，例如通过循环信念传播优化（loopy belief propagation,LBP）。[7]和[18]将图模型中的平滑约束集成到端到端的学习分割网络。由于LBP需要更多迭代才的带最优结果，因此并不合适作为展开的网络层，我们使用了用于[8]中的离散优快速平均场推断（两次迭代），在[8中使用了5次迭代。它由两个交替的步骤组成：标签兼容变换（独立地作用域空间控制点）和基于滤波器的消息传递（使用平均池化层实现，步长为1）。

正如[3]中所指出的那样，稠密位移空间的扩散正则化可以使用以3D位移偏移为根的抛物线的较低包络来计算，其高度等于不同项的和与之前的平均场推断迭代。这个较低的包络线不能直接微分，但是我们可以得到非常精确的近似值，首先使用最小池化（stride=1）来找到局部极小值然后使用两个平均代价池操作average pooling operation(stride = 1)来提供二次平滑。正如图1中蓝色快所示，我们的方法新正则化部分包括左右与3个位移维度（min-convolution）的最小和平均池层，以及作用于3个空间维（mean-field inference）的平均滤波。每个操作、放缩和偏差因素a1-a6之前介绍和优化功能层在端到端一起训练

概率变换损失和标签监督：我们可以进一步利用位移采样的概率性质，并根据非局部加权[11]的方法指定我们的监督监督标签损失项。即我们首先否定正则部分的输出（缩放）到伪概率使用softmax计算位移。然后，在相同空间移位位置对移动分割的one-hot表示进行采样，并将这些向量乘以估计概率，一计算作为MSE的真实值（one-hot）分割的标签损失。3D位移场的连续值是将概率估计与位移标签相乘的加权平均值，然后对图像分辨率进行三线性插值得到。扩散正则惩罚在位移场的所有的3维空间梯度是用户定义用于平衡平滑变换（较低的Jacobians的标准偏差）和准确结构调整。

3 实验验证

为了证明我们的方法能够在不同病人腹部捕捉到非常大的形变，我们对CIS-CERAL3的10个经过对比增强的3D CT扫描的训练数据，进行了3倍交叉验证实验，每9个解刨结构手动分割：肝脏liver、脾脏spleen、胰腺pancreas、胆囊gallbladder、膀胱unary bladder，右肾right kidney, 左肾left kidney,右腰大肌 right psoas major muscle(psoas)和左腰大肌left psoas。将图像重采样到各向同性体素大小为1.5mm^3，尺寸为233x168x286个体素，不需要人工预对齐。

我们将概率密度置换网络(probabilistic dense displacement network ,pdd-net)与两种传统算法NiftyReg和deeds进行比较，这两种算法在[16]的患者间腹部CT配准研究表现最佳，但DLIR尚未处理这一任务。NiftyReg与互信息mutual information和5级多分辨率方案一起用于捕获大的形变运行时间40~50s。Deed考虑单尺度密度的移位空间（大概需要4-6s），然后扩展到离散优化的三等级（25-35秒的运行时）。下一步，我们在数据上训练了弱监督DLIR方法Label-Reg在我们数据上（每个折叠在24小时以上）。为了将内存需求降低到32G以下，分辨率降低到2.2mm基本通道数量减半到16。进一步践行了小的调整，以优化病人之间的训练。我们实现了FlowNet-c的三维扩展在带有Obelisk特征提取的pytorch中，一个密集相关层a dense correlation layer和一个具有输入通道的正则化网络由5个batch-norma层高的PReLU。它有200万个权值和输出一个(非概率)三维位移场。为了获取更有意义的结果，需要在Obelisk层的中间输出加入语义分割损失。我们提出的方法使用相同的特征学习部分（200k参数），但现在使用min-convolution，mean-field推理（没有语义指导）和non-local label loss只增加6个可训练的权重（不是2百万个）。在消融研究中分析了这三种选择的影响，并考虑用手工自相似上下文特征代替Obelisk特征学习。我们使用扩散的正则项权重用于控制32^3的网格大小和fixed scans的整个仿射增强和训练我们Adma网络（学习率0.01）在1500次迭代大约90分钟GPU显存16G。我们实现了一个实例梯度下降优化器，优化了前馈预测。[1]也使用了这个思想，但在我们的例子中，它快了100倍（0.24s比24s），因为我们可以直接操作预先计算的位移概率，并且不需要通过网络进行迭代反向传播。

表1。定量比较交叉验证的VISCERAL解刨3个数据集的10个扫描，基于24个组合的测试扫描没有在训练集上看到（数字表示Dice得分）。我们的pdd-net性能优于其他DLIR方法Label-Reg和FlowNet-C，比传统方法（NiftyReg和deeds）的优势在15%。我们剥离研究显示1）用于学习Obelisk特征与手动自相似上下文（SSC）描述符。2）使用平均场推断。3）使用心得非局部标签损失。另外一个快速实例级优化实现了pdd-net+inst。 FlowNet-C是我们[2]的3D扩展，具有Obelisk特征和可训练的正则化网络。比较配准前的Dice平均为30%。

4、结果和讨论

pdd-net的推理时间只是0.57秒，产生的位移场似是而非，Jacobian行列式标准差为0.4小于1%折叠体素(负雅可比矩阵)。表1显示了Dice平均分数交叉验证了24组配准。其中为使用任何标记的训练扫描进行任何评估的测试配准。我们的方法比Label-Reg和FlowNet-C两种DL方法，以大约15%个点，只有30%的初始对齐的病人间挑战达到56.7%的Dice系数。这比一个网格级别的传统离散配准deeds要好10%。我们的实例优化（每扫描对）需要0.24秒，较少折叠（少于0.6%）并且进一步将准确率提高到了58.4%，这是高于传统的多级配准deeds和NiftyReg。

比较deeds+SSC和一个网格级pdd+SSC，到我们pdd+SSC的变体。它使用相同的自相似性的特点，只有适应的正则部分的α参数，我们得到一个相思精确和形变复杂性。这表明所提出的最小卷积和两个平均场推断的正则化层几乎可以匹配全序列的能力[5]中MRF优化。使用弱监督学习的特点，结果增加20%以上的Dice。非局部损失项和我们实例微调，分别贡献了5%和2%进一步收益的Dice重叠。考虑到具有更多可训练权值的无约束FlowNet-C的效果更差，或者我们的变体只是用小卷积而不是进行空间域滤波，均值域推断的重要性是显而易见。的。我们实现了一个比常规配准更鲁棒的对齐质量（Dice更低的标准差）。可视化配准示例如图2所示，并作为补充材料中的表面呈现视频文件。

5、结论

我们的新型pdd-net将概率稠密位移与可微分的平均场正则化相结合，在7个更大的解剖病人间腹部CT配准，达到了70%Dice以上的one-to-one精度。它比之前的基于深度学习的图像配准方法（Label-Reg和FlowNet-C），要高出15%，并且可以通过少量的标记臊面进行鲁棒性训练。它弥补了DLIR（在小训练数据集)与传统方法的质量差距，例如NiftyReg和deeds，同时快了(0.5秒)。我们的概念提供了一个新的潜在方向，使得DLIR在图像引导介入、诊断和基于atlas形状分析中的运用，超越了目前使用缺乏几何可解释性的像素分割网络。未来的工作可以通过使用多个校准阶段和控制点的更自适应采样来获得进一步的收益。对具有附加评价指标(表面距离)的较大数据集进行更详细的评价，可以对该方法的优缺点提供更多的见解。