【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

本文是论文《Cycle-Consistent Training for Reducing Negative Jacobian Determinant in Deep Registration Networks》的阅读笔记。

微分同胚

配准任务中使用的变换(形变场)理想情况下应该是微分同胚的,这样图像的拓扑性在转换的时候就不会发生变化。微分同胚即是可逆和光滑的,现有的配准方法通常不能从理论上保证形变场的平滑性和可逆性。

【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

之前的无监督的深度学习方法是通过使用局部平滑约束来直接使形变场保持平滑,或者对网络结构进行设计来保证形变场平滑。文章通过研究可能的训练机制/任务,这些机制/任务将帮助网络避免使用负雅克比行列式预测变换并产生更平滑的变形。

【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

上图是配准的大体框架(VoxelMorph就是采用以上结构),其中采样器,即空间转换网络(STN)由变形单元(deformation unit)和采样单元(sampling unit)组成,假设源域图像和目标域图像分别为xyx,y,则变形单元会产生一个静态位移场uu,那么就可以根据位移场uu,使用插值的方法从源图像xx得到采样后的图像y~=x(Id+u)\tilde{y}=x(Id+u),其中IdId表示恒等映射。

微分同胚ϕ\phi对图像xx的操作被近似为ϕx=xϕ1x(Id+u)\phi \cdot x=x \circ \phi^{-1} \approx x(I d+\mathbf{u}),位移场uu的平滑性通常通过对其导数DuDu添加正则化来保证。比如VoxelMorph网络的损失函数为:CC(y,y~)+λDul2C C(y, \tilde{y})+\lambda\|D \mathbf{u}\|_{l_{2}},其中CCCC表示互相关,λ\lambda是控制正则项的权重。
【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

形变场的雅克比行列式,即det(Dϕ1)det(Id+Du)\operatorname{det}\left(D \phi^{-1}\right) \approx \operatorname{det}(\mathrm{Id}+D \mathbf{u})很有可能在多个位置都是负的,即使增大λ\lambda的值时,这种重叠问题依然可能会存在。此外,超参λ\lambda的确定也是比较麻烦的,比较大的λ\lambda值会导致较小的位移场,而降低模型的精度。如上图所示,分别是λ=1,2,4\lambda=1,2,4时形变场的折叠情况。

Prob-VoxelMorph采用一个重参数的技巧,并插入集成层(integration layer)来让模型产生平滑的形变场。循环一致性是通过为模型增加一个额外的任务来间接的解决形变场的平滑性问题。

循环一致性

不同于以往的模型,文本不是通过增加一个特定的损失项来防止形变场重叠,而是通过研究有可能的训练机制/任务来实现。循环一致性的想法不会改变backbone网络的结构,而是通过增加一个从已经预测的图像中恢复原始图像的任务来实现正则化。
【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

循环一致性具体来说就是,对于图像xyx,y,先将yy作为固定图像,将xx配准到yy,得到配准图像y~\tilde y,该变换的位移场为uxy~u_{x\rightarrow\tilde y};然后再以xx作为固定图像,将y~\tilde y配准到图像xx,该变换的位移场为uy~x~u_{\tilde y\rightarrow\tilde x}。上图是循环一致性配准的流程示意图。整个网络的循环一致性损失为:
CC(y,y~)+λDuxy~l2+CC(x,x~)+λDuy~x~l2 C C(y, \tilde{y})+\lambda\left\|D \mathbf{u}_{x \rightarrow \tilde{y}}\right\|_{l_{2}}+C C(x, \tilde{x})+\lambda\left\|D \mathbf{u}_{\tilde{y} \rightarrow \tilde{x}}\right\|_{l_{2}}
这种设计直接解决了网络的可逆性,循环一致性任务同时也间接保证了网络可以学习到一个平滑的形变场,即帮助网络排除那些非循环一致的变换。此外这种设计并不会增加模型的参数。
【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

循环一致性的想法虽然跟双向配准有些相似,但是还是有很大的不同。在循环一致性配准中,目标图像从没发生变形。双向配准是既使用(x,y)(x,y)又使用(y,x)(y,x)作为输入,而循环一致性配准只使用(x,y)(x,y)作为输入。上表是两种方式的对比。

实验

实验使用的是MindBoggle101数据集,该数据集包括三个子数据集:

  • NKI-RS-22
  • NKI-TRT-20
  • OASIS-TRT-20

每个图像都是182×218×182182\times218\times182大小的,然后裁剪到144×180×144144\times180\times144大小,这些图像已经对齐到MNI152空间,还需要将其灰度值归一化到[0, 1]的范围。原始的VoxelMorph将作为本文模型的backbone网络,实验具体设置为λ=1\lambda=1,epoch=10,选用学习率为1e41e^{-4}的Adam优化器,使用三折交叉验证。使用Dice值作为评价指标,其表达式如下:
Dice((ϕLxi),Lyi)=2(ϕLxi)LyiϕLxi+Lyi \operatorname{Dice}\left(\left(\phi \cdot L_{x}^{i}\right), L_{y}^{i}\right)=\frac{2\left|\left(\phi \cdot L_{x}^{i}\right) \cap L_{y}^{i}\right|}{\left|\phi \cdot L_{x}^{i}\right|+\left|L_{y}^{i}\right|}
其中LxiL^i_x是带有下标的标签,Dice值越大说明配准效果越好。

使用OASIS-TRT-20作为测试集。形变场的重叠是通过检查负雅可比行列式发生的位置来获得的。设PP定义为所有体素VV上Jacobian行列式为负的体素位置的百分比,即:
P:=δ(det(Dϕ1)<0)V \mathcal{P}:=\frac{\sum \delta\left(\operatorname{det}\left(D \phi^{-1}\right)<0\right)}{V}
理想情况下,PP的值越小越好。
【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

上图表示三种不同方法的雅克比行列式和采样网格情况,其中红色的部分表示行列式为负数的位置。在baseline模型中超过90%的不满意的位置都被消除了。所以有两种不同的方式可以减少形变场的重叠位置,同时保证配准的精度,一种是像Prob-VoxelMorph一样的直接方式,另一种是像循环一致性配准一样的间接方式。
【论文笔记】用循环一致性避免形变场重叠的医学图像配准网络

上表是几种不同配准方法的PP和Dice的均值和标准差对比。