【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

本文是论文《CycleMorph: Cycle Consistent Unsupervised Deformable Image Registration》的阅读笔记。

现有的既又深度学习的配准方法通常不能保证拓扑保持(preservation of topology),为了解决这个问题,文章提出了一个名为CycleMorph的模型,使用循环一致性来提供一个隐式的正则来保留拓扑结构。此外,该模型可以实现2D或3D图像以及多尺度图像的配准。

一、相关工作

1. CycleMorph使用循环一致性来

在该模型中训练两个卷积神经网络GXG_XGYG_YGXG_X产生从源域图像到目标域图像的正向形变场ϕXY\phi_{XY}GYG_Y产生从目标域图像到源域图像的逆向形变场ϕYX\phi_{YX}

为了解决3D图像太大而导致显存放不下的情况,采用了先对下采样的图像进行粗糙的大变形,然后再对局部的patch进行变形。

2. 微分同胚图像配准

在经典的变分图像配准方法中,配准问题的能量方程可以表示为:
L(X,Y,ϕ)=Lsim(T(X,ϕ),Y)+Lreg(ϕ) \mathcal{L}(X, Y, \phi)=\mathcal{L}_{s i m}(\mathcal{T}(X, \phi), Y)+\mathcal{L}_{r e g}(\phi)
其中,X,YX,Y分别表示浮动图像和固定图像,ϕ\phi表示位移向量场,T\mathcal T表示根据形变场ϕ\phiXX变形到YY的变换函数(空间变换层),Lsim\mathcal L_{sim}是衡量图像相似性的相似项,Lreg\mathcal{L}_{reg}是使得形变场平滑的正则项。

微分同胚图像配准为向量场ϕ\phi施加了一个约束,使得可变形映射是微分同胚。微分同胚变形保证了两个图像之间的变形是连续、可微、拓扑保持的。微分同胚图像配准算法的例子有LDDMM和SyN。

无监督的图像配准方法通常没有为一致性增加一个约束,从而可能会产生形变场重叠的问题。

3. 一致性图像配准

当参数有限时,微分同胚图像配准算法中的形变场通常是离散的。因此,估计的形变F:XYF: X \mapsto Y和估计的形变R:YXR: Y \mapsto X的逆是不相等的。在一致性图像配准中,这个问题通过逆一致性RF1R\simeq F^{-1}而减轻了。

二、CycleMorph

【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是CycleMorph网络的结构示意图。

1. 损失函数

模型通过优化以下损失函数来训练:
minGX,GYL(X,Y,GX,GY) \min _{G_{X}, G_{Y}} \mathcal{L}\left(X, Y, G_{X}, G_{Y}\right)
其中
L(X,Y,GX,GY)=Lregist(X,Y,GX)+Lregist(Y,X,GY)+αLcycle(X,Y,GX,GY)+βLidentity(X,Y,GX,GY) \begin{aligned} \mathcal{L}\left(X, Y, G_{X}, G_{Y}\right)=& \mathcal{L}_{r e g i s t}\left(X, Y, G_{X}\right) \\ &+\mathcal{L}_{r e g i s t}\left(Y, X, G_{Y}\right) \\ &+\alpha \mathcal{L}_{c y c l e}\left(X, Y, G_{X}, G_{Y}\right) \\ &+\beta \mathcal{L}_{i d e n t i t y}\left(X, Y, G_{X}, G_{Y}\right) \end{aligned}
【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是网络中三个损失的示意图。Lregist\mathcal{L}_{regist}是配准损失,它包括图像相似性损失和形变场的正则项损失两部分,图像相似性采用的是局部互相关,正则项损失采用的是L2L_2正则损失,所以
Lregist(X,Y,GX)=(T(X,ϕXY)Y)+λϕXY2 \begin{array}{l} \mathcal{L}_{\text {regist}}\left(X, Y, G_{X}\right) \\ \quad=-\left(\mathcal{T}\left(X, \phi_{X Y}\right) \otimes Y\right)+\lambda \sum\left\|\nabla \phi_{X Y}\right\|^{2} \end{array}
其中λ\lambda是超参数,\otimes表示局部互相关,其定义为:
AB=vΩ(vi(A(vi)Aˉ(v))(B(vi)Bˉ(v)))2(vi(A(vi)Aˉ(v))2)(vi(B(vi)Bˉ(v))2) A \otimes B=\sum_{\mathbf{v} \in \Omega} \frac{\left(\sum_{\mathbf{v}_{i}}\left(A\left(\mathbf{v}_{i}\right)-\bar{A}(\mathbf{v})\right)\left(B\left(\mathbf{v}_{i}\right)-\bar{B}(\mathbf{v})\right)\right)^{2}}{\left(\sum_{\mathbf{v}_{i}}\left(A\left(\mathbf{v}_{i}\right)-\bar{A}(\mathbf{v})\right)^{2}\right)\left(\sum_{\mathbf{v}_{i}}\left(B\left(\mathbf{v}_{i}\right)-\bar{B}(\mathbf{v})\right)^{2}\right)}
其中,Ω\Omega表示整个3D图像,Aˉ(v)\bar{A}(v)Bˉ(v)\bar{B}(v)表示图像A(v)A(v)B(v)B(v)的局部均值,viv_i迭代在体素vv周围大小为w×w×ww\times w\times w(2D图像时为w×ww\times w)的范围迭代计算,其中w=9w=9

Lcycle\mathcal{L}_{cycle}是循环一致性损失,源域图像XX先通过GXG_X被配准得到Y^\hat{Y},再通过GYG_Y配准得到X~\tilde{X},计算XXX~\tilde{X}之间的循环一致性损失;同理目标域图像YY先通过GYG_Y被配准得到X^\hat{X},再通过GXG_X配准得到Y~\tilde{Y},计算YYY~\tilde{Y}之间的循环一致性损失。上述过程可以表示为:
(X,Y)(T(Y^,ϕ^YX),T(X^,ϕ^XY)) (X, Y) \simeq\left(\mathcal{T}\left(\hat{Y}, \hat{\phi}_{Y X}\right), \mathcal{T}\left(\hat{X}, \hat{\phi}_{X Y}\right)\right)
其中
(Y^,X^):=(T(X,ϕXY),T(Y,ϕYX)) (\hat{Y}, \hat{X}):=\left(\mathcal{T}\left(X, \phi_{X Y}\right), \mathcal{T}\left(Y, \phi_{Y X}\right)\right)
所以,循环一致性损失可以通过下式计算:
Lcycle(X,Y,GX,GY)=T(Y^,ϕ^YX)X1+T(X^,ϕ^XY)Y1 \begin{array}{l} \mathcal{L}_{\text {cycle}}\left(X, Y, G_{X}, G_{Y}\right) \\ \quad=\left\|\mathcal{T}\left(\hat{Y}, \hat{\phi}_{Y X}\right)-X\right\|_{1}+\left\|\mathcal{T}\left(\hat{X}, \hat{\phi}_{X Y}\right)-Y\right\|_{1} \end{array}
其中,1\|\cdot\|_1表示L1L_1范式。

Lidentity\mathcal{L}_{identity}是恒等损失表示当固定图像和浮动图像为同一个时,两者之间不会有任何变形。因此该损失避免了不必要的变形,使得形变场在静态区域的预测更加稳定。该损失可以表示为:
Lidentity(X,Y,GX,GY)=(T(Y,GX(Y,Y))Y)(T(X,GY(X,X))X) \begin{array}{l} \mathcal{L}_{\text {identity}}\left(X, Y, G_{X}, G_{Y}\right) \\ =-\left(\mathcal{T}\left(Y, G_{X}(Y, Y)\right) \otimes Y\right)-\left(\mathcal{T}\left(X, G_{Y}(X, X)\right) \otimes X\right) \end{array}
其中\otimes表示局部互相关。

2. 空间变换层

对于3D图像来说,使用三线性插值,可以表示为:
T(X,ϕ)=qN(p+ϕ(p))X(q)d{i,j,k}(1pd+ϕ(pd)pd) \begin{array}{l} \mathcal{T}(X, \phi) \\ =\sum_{q \in \mathcal{N}(p+\phi(p))} X(q) \prod_{d \in\{i, j, k\}}\left(1-\left|p_{d}+\phi\left(p_{d}\right)-p_{d}\right|\right) \end{array}
其中,pp表示像素下标,N(p+ϕ(p))\mathcal{N}(p+\phi(p))表示在p+ϕ(p)p+\phi(p)周围的8像素的立方体邻居,dd是3D图像空间的三个方向。当为2D图像配准时,采用的是双线性插值。

3. 多尺度配准

【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是多尺度配准的示意图,训练时多尺度配准采用先做全局配准,后做局部配准。全局图像配准模型首先在下采样的图像对上进行训练,然后通过对形变场进行上采样得到全尺寸的变形后的图像,然后对变形后图像的patch进行局部图像配准的训练。

在测试阶段,由于要进行两次线性插值,所以两个配准网络会导致误差的累计,从而降低配准的精度。所以采取了全局配准和局部配准先得到各自的形变场,然后再将形变场相加,只做一次线性插值。

三、实验

1. 数据集

三个任务三个数据集,2D配准任务使用的是面部表情图像Radboud Faces Database(RaFD)数据集,3D配准任务使用的是OASIS-3脑部MRI数据集,多尺度配准使用的是肝脏CT数据集。不同任务所采用的参数就不再赘述。

实验采用的baseline有Elastix,SyN,VoxelMorph和MS-DIRNet。

2. 评价指标

使用形变场的雅克比行列式中非负数的百分比来衡量形变场的重叠程度:
Jϕ(v)=ϕ(v)0 \left|J_{\phi}(\mathbf{v})\right|=|\nabla \phi(\mathbf{v})| \leq 0
其中vv表示体素位置,|\cdot|表示矩阵的行列式,当雅克比矩阵的行列式的值都是正数时,形变场是微分同胚的。

在面部表情数据集中有ground-truth标签,衡量其配准后的归一化均方误差损失(NMSE)和结构相似性损失(SSIM);在脑部MR数据集中有脑部结构的分割标签,衡量其配准后的DICE值,其表达式如下:
Dice(A,B)=2TP2TP+FP+FN \operatorname{Dice}(A, B)=\frac{2 T P}{2 T P+F P+F N}
其中TP,FP,FNTP,FP,FN分别表示正阳例、负阳例、负阴例像素的个数。

在肝脏CT数据集中有地标点信息,衡量其在配准后的目标配准损失(TRE),TRE可以通过平均欧几里得距离来计算:
TRE(A,B)=1Ni=1Naibi T R E(A, B)=\frac{1}{N} \sum_{i=1}^{N}\left\|a_{i}-b_{i}\right\|
其中NN是地标点的数目,aibia_i,b_i是浮动图像和固定图像中第ii个地标坐标向量。

3. 实验结果

【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是在面部表情数据集上的实验结果。
【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是在脑部MR数据上的配准结果。

【论文笔记】CycleMorph:使用循环一致性的2D/3D/多尺度微分同胚图像配准模型

上图是不同配准方法在脑部MR数据集上的配准结果的对比表。