医学图像配准 | SYMnet 对称微分同胚配准CNN（SOTA） | CVPR2020

机器学习炼丹术
1,139 阅读4分钟
  • 文章转自：微信公众号「机器学习炼丹术」
  • 作者：炼丹兄（已授权）
  • 作者联系方式：微信cyx645016617
  • 论文名称：‘Fast Symmetric Diffeomorphic Image Registration with Convolutional Neural Networks’
  • 论文链接：arxiv.org/abs/2003.09…

0 综述

1 微分同胚回顾

ϕ(1)\phi^{(1)}是表示时间间隔为1的形变场，上图中的x表示形变场，然后形变场加上形变场的速度，就是时间间隔更长的形变场。

根据李代数群论，得到的结论是： ϕ(1/2)=ϕ(1/4)∘ϕ(1/4)\phi^{(1/2)}=\phi^{(1/4)} \circ \phi^{(1/4)},我们可以从时间间隔为1/4的变形场，推出时间间隔1/2的变形场。

从上一篇文章的代码中我们也可以知道，论文中关于这部分的实现是下面的逻辑，以SYMnet中的T=7为例，是吧每一个时间间隔都划分成了272^7小段。

  • 把模型输入的形变场，除以272^7，这样就是所谓的ϕ1/27\phi^{1/2^7}
  • 通过ϕ1/27∘ϕ1/27\phi^{1/2^7} \circ \phi^{1/2^7}得到ϕ1/26\phi^{1/2^6}
  • 通过ϕ1/26∘ϕ1/26\phi^{1/2^6} \circ \phi^{1/2^6}得到ϕ1/25\phi^{1/2^5}
  • 最后可以得到ϕ(1)\phi^{(1)}

2 模型结构

这个模型的思想是，X和Y两个图片配准，本来是从X到Y的，但是现在我们需要找到X和Y的一个中间态Z，让X配准到Z，然后让Y也配准到Z。

图中的ϕXY(1)\phi_{XY}^{(1)}是从X到Y的形变场。

最终推理部分，我们肯定是需要从X配准到Y的。所以我们先用ϕXY0.5\phi_{XY}^{0.5}让X配准到中间态Z，然后再用ϕYX(−0.5)\phi_{YX}^{(-0.5)}让中间态Z逆向配准到Z。

2.1 FCN

特征提取部分依然是使用U-net（原文中是类似U-net的FCN）：

  • 网络的输入依然是X和Y两个图片拼接起来的，2通道图片；
  • 论文中说，在模型的最后，两层卷积层为5的卷积层被用来生成两个速度场，vXY,vYXv_{XY},v_{YX}
  • 在后面跟着一个softsign activation x1+∣x∣\frac{x}{1+|x|}
  • 然后乘上一个常数c，让速度的范围在[-c,c]，论文中给定的c为100，这样对于大的形变是有效的；
  • 除了output卷积层，每一个卷积层后面都跟着ReLU激活层。
  • 根据微分同胚的那个推论，假如设置的是T=7的话，那么我们最后只需要ϕ(0.5)\phi^{(0.5)}就可以了。

3 损失函数

  • 最主要的损失就是NCC，对于这种图像配准任务来说，这个损失一般都是有的。论文中用的是NCC，也可以使用MSE等

论文中SYMnet模型的loss有：

  • LsimL_{sim}衡量相似性的损失；
  • LJdetL_{Jdet}利用雅各比行列式来代替梯度平滑损失。

3.1 similarity loss

包含两个部分： Lsim=Lmean+LpairL_{sim} = L_{mean} + L_{pair}：

  • Lmean=−NCC(X(ϕXY(0.5)),Y(ϕYX(0.5)))L_{mean}=-NCC(X(\phi_{XY}^{(0.5)}),Y(\phi_{YX}^{(0.5)}))。这个很好理解，就是希望X转换成中间态的Z和Y转换成的中间态Z是相同的；
  • Lpari=−NCC(X(ϕXY(1)),Y)−NCC(Y(ϕYX(1)),X)L_{pari} = -NCC(X(\phi_{XY}^{(1)}),Y)-NCC(Y(\phi_{YX}^{(1)}),X).这个也好理解，就是希望X转换成的Y和Y转换的X和真实的Y和X是相同的；

其中值得一提的是:

  • ϕXY(1)=ϕXY(0.5)∘ϕYX(−0.5)\phi_{XY}^{(1)} = \phi_{XY}^{(0.5)} \circ \phi_{YX}^{(-0.5)}
  • ϕYX(1)=ϕYX(0.5)∘ϕXY(−0.5)\phi_{YX}^{(1)} = \phi_{YX}^{(0.5)} \circ \phi_{XY}^{(-0.5)}

3.2 雅各比行列式损失

作者提出了雅各比行列式损失来代替voxelmorph的smooth grad loss。

这个损失是代替之前的grad smooth loss，更注重局部方向的一致性。使用现有的方法，例如L1或者L2正则对变形场上的梯度进行约束，这种全局正则化会大大降低配准的精度

论文中提出雅各比行列式方法，对估计的变形场施加局部方向的一致性约束

其中的N表示变形场的元素数量，σ\sigma是ReLU激活函数，Jϕ(p)J_{\phi}(p)表示变形场在p位置的雅各比行列式。行列式的定义如下：

代码如下：

def JacboianDet(y_pred, sample_grid):
    J = y_pred + sample_grid
    dy = J[:, 1:, :-1, :-1, :] - J[:, :-1, :-1, :-1, :]
    dx = J[:, :-1, 1:, :-1, :] - J[:, :-1, :-1, :-1, :]
    dz = J[:, :-1, :-1, 1:, :] - J[:, :-1, :-1, :-1, :]

    Jdet0 = dx[:,:,:,:,0] * (dy[:,:,:,:,1] * dz[:,:,:,:,2] - dy[:,:,:,:,2] * dz[:,:,:,:,1])
    Jdet1 = dx[:,:,:,:,1] * (dy[:,:,:,:,0] * dz[:,:,:,:,2] - dy[:,:,:,:,2] * dz[:,:,:,:,0])
    Jdet2 = dx[:,:,:,:,2] * (dy[:,:,:,:,0] * dz[:,:,:,:,1] - dy[:,:,:,:,1] * dz[:,:,:,:,0])

    Jdet = Jdet0 - Jdet1 + Jdet2

    return Jdet


def neg_Jdet_loss(y_pred, sample_grid):
    neg_Jdet = -1.0 * JacboianDet(y_pred, sample_grid)
    selected_neg_Jdet = F.relu(neg_Jdet)

    return torch.mean(selected_neg_Jdet)

It is worth noting that the proposed selective Ja- cobian determinant regularization loss means not to replace the global regularizer. Instead, we utilize both regulariza- tion loss functions in our method to produce smooth and topology-preservation transformations while alleviating the tradeoff between smoothness and registration accuracy.

作者强调，这个雅各比行列式损失，并没有替代之前的全局正则化损失，而是两者并重。

所以下面是L2的梯度正则化损失：

然后为了还增加了magnitude损失，就是确保vxyv_{xy}和xyxx_{yx}有着相同的量级：

所以总的来说，这个SYMnet有四个部分的损失：

4 结果

不出意外，结果也是实现了最好的水准。

文章主打的是雅各比损失，所以专门做了雅各比损失权重不同对最终结果的影响

有意思的是，虽然作者说∣Jϕ∣越小越好，但是我看到的是DSC的效果随着雅各比权重的降低而提升哈哈。似乎没有这个损失约束，效果是最好的哈哈|J_{\phi}|越小越好，但是我看到的是DSC的效果随着雅各比权重的降低而提升哈哈。似乎没有这个损失约束，效果是最好的哈哈

avatar
文章
阅读
粉丝