【论文笔记】CIRNet:基于CycleGAN的无监督循环配准模型

本文是论文《UNSUPERVISED THREE-DIMENSIONAL IMAGE REGISTRATION USING A CYCLE CONVOLUTIONAL NEURAL NETWORK》的阅读笔记。

文章提出了一个无监督的循环图像配准网络CIRNet,它由两个配准模块组成,两者具有相同的结构,并且共享参数。此外加入了循环一致性损失来为确保预测形变场的精度而提供额外的约束。文章在4D(3D+t)心脏CT数据集和MRI数据集上分别做了实验,验证了模型的有效性。简单来说就是将CycleGAN引入到了VoxelMorph模型中。

一、配准问题介绍

有监督的配准方法容易收到有监督信息中存在的误差和错误的影响,基于patch的配准方法训练时更快,但不能充分利用全局信息。

配准问题可以简化为一个优化问题,即最小化以下能量函数 P P
P = S ( M ( ϕ ) , F ) + Reg ( ϕ ) P=S(M(\phi), F)+\operatorname{Reg}(\phi)
其中, M M F F 分别表示浮动图像和固定图像, ϕ \phi 表示形变场, M ( ϕ ) M(\phi) 表示根据形变场变形后的浮动图像,函数 S ( x , y ) S(x,y) 表示两幅图像之间的不相似性, R e g ( ϕ ) Reg(\phi) 表示形变场的平滑正则化项。

二、网络结构

从下图可以看出,CIRNet由两个级联的配准模块 R 1 R 2 R_1,R_2 组成,两者具有相同的结构,并且共享参数,这样模型参数不会增加很多。

上图中, A A B B 分别表示浮动图像和固定图像, A ( ϕ 1 ) A(\phi_1) 是根据形变场 ϕ 1 \phi_1 变形后的浮动图像, A ( ϕ 1 ϕ 2 ) A(\phi_1\circ\phi_2) 是先根据 ϕ 1 \phi_1 变形再根据 ϕ 2 \phi_2 变形后的浮动图像, L c y c L_{cyc} 是循环一致性损失。配准模块 R 1 R_1 是将浮动图像 A A 配准到固定图像 B B ,得并使得 A ( ϕ 1 ) A(\phi_1) B B 尽可能的像;配准模块 R 2 R_2 是将变形后的浮动图像 A ( ϕ 1 ) A(\phi_1) 配准到固定图像 A A ,得并使得 A ( ϕ 1 ϕ 2 ) A(\phi_1\circ\phi_2) A A 尽可能的像。

CIRNet是在VoxelMorph的基础上加以改进的,所以每个配准模块的结构和VoxelMorph相同,如下图。

三、损失函数

网络的损失有三个: L N C C L c y c L r e g L_{NCC}、L_{cyc}、L_{reg}

归一化的互相关 L N C C L_{NCC} ,它用来衡量图像 A ( ϕ 1 ) A(\phi_1) 和图像 B B 之间的相似性,它是通过大小为 9 × 9 × 9 9\times9\times9 的滑动窗口来计算的,当 L N C C L_{NCC} 接近于0时,说明两张图像不相关,反之接近于1时说明两张图像正相关。

循环一致性损失 L c y c L_{cyc} 可以让图像 A ( ϕ 1 ϕ 2 ) A(\phi_1\circ\phi_2) 和图像 A A 尽可能的相似。此外考虑到两幅图像还应该在灰度值分布上相同,所以在循环一致性损失中加入了 L 2 L_2 正则项。其表达式如下:
L c y c ( A ( ϕ 1 ϕ 2 ) , A ) = 1 N i Ω ( A ( ϕ 1 ϕ 2 ) ( i ) A ( i ) ) 2 L_{c y c}\left(A\left(\phi_{1} \circ \phi_{2}\right), A\right)=\frac{1}{N} \sum_{i \in \Omega}\left(A\left(\phi_{1} \circ \phi_{2}\right)(i)-A(i)\right)^{2}
其中, N N 是体素的数量, Ω \Omega 表示图像中的所有体素。

形变场的平滑正则项损失 L r e g L_{reg} 定义如下:
L r e g ( ϕ ) = i Ω ϕ ( i ) 2 L_{r e g}(\phi)=\sum_{i \in \Omega} \nabla \phi(i)^{2}
综上,配准模块 R 1 R 2 R_1、R_2 的总损失为:
L R 1 ( A , B ) = L N C C ( A ( ϕ 1 ) , B ) + α L r e g ( ϕ 1 ) + β L c y c ( A ( ϕ 1 ϕ 2 ) , A ) L R 2 ( A ( ϕ 1 ) , A ) = L N C C ( A ( ϕ 1 ϕ 2 ) , A ) + α L r e g ( ϕ 2 ) + β L c y c ( A ( ϕ 1 ϕ 2 ) , A ) \begin{aligned} L_{R 1}(A, B)=-& L_{N C C}\left(A\left(\phi_{1}\right), B\right) \\ &+\alpha L_{r e g}\left(\phi_{1}\right)+\beta L_{c y c}\left(A\left(\phi_{1} \circ \phi_{2}\right), A\right) \\ L_{R 2}\left(A\left(\phi_{1}\right), A\right) &=-L_{N C C}\left(A\left(\phi_{1} \circ \phi_{2}\right), A\right) \\ &+\alpha L_{r e g}\left(\phi_{2}\right)+\beta L_{c y c}\left(A\left(\phi_{1} \circ \phi_{2}\right), A\right) \end{aligned}
其中 α β \alpha、\beta 是平衡因子。

四、实验

第一个数据集是ACDC数据集,是一个4D的心脏电影MRI数据集,第二个数据是心脏的CTA图像。训练时使用Adam作为优化器,学习率在前100个epoch为0.0001,没50个epoch线性减少。batch size在第一个数据集为8,第二个数据集为1。 α = 1 β = 5 \alpha=1,\beta=5 。使用的baseline是传统的B样条非刚性配准算法Elastix和基于学习的配准模型VoxelMorph,评价指标选用的是DSC(Dice相似性系数)。

下图是各个配准方法的对比结果。
在这里插入图片描述
下图是评价指标和运行时间的对比,可以发现CIRNet取得了最好的效果。

在这里插入图片描述