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

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

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

一、配准问题介绍

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

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

二、网络结构

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

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

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

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

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

三、损失函数

网络的损失有三个:LNCCLcycLregL_{NCC}、L_{cyc}、L_{reg}

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

循环一致性损失LcycL_{cyc}可以让图像A(ϕ1ϕ2)A(\phi_1\circ\phi_2)和图像AA尽可能的相似。此外考虑到两幅图像还应该在灰度值分布上相同,所以在循环一致性损失中加入了L2L_2正则项。其表达式如下:
Lcyc(A(ϕ1ϕ2),A)=1NiΩ(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}
其中,NN是体素的数量,Ω\Omega表示图像中的所有体素。

形变场的平滑正则项损失LregL_{reg}定义如下:
Lreg(ϕ)=iΩϕ(i)2 L_{r e g}(\phi)=\sum_{i \in \Omega} \nabla \phi(i)^{2}
综上,配准模块R1R2R_1、R_2的总损失为:
LR1(A,B)=LNCC(A(ϕ1),B)+αLreg(ϕ1)+βLcyc(A(ϕ1ϕ2),A)LR2(A(ϕ1),A)=LNCC(A(ϕ1ϕ2),A)+αLreg(ϕ2)+βLcyc(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:基于CycleGAN的无监督循环配准模型
下图是评价指标和运行时间的对比,可以发现CIRNet取得了最好的效果。

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