Deep Learning-Based Concurrent Brain Registration and Tumor Segmentation

Estienne, T. et al. “Deep Learning-Based Concurrent Brain Registration and Tumor Segmentation.” Frontiers in Computational Neuroscience 14 (2020): n. pag.

Abstract

Image registration and segmentation are the two most studied problems in medical image analysis. Deep learning algorithms have recently gained a lot of attention due to their success and state-of-the-art results in variety of problems and communities. In this paper, we propose a novel, efficient, and multi-task algorithm that addresses the problems of image registration and brain tumor segmentation jointly. Our method exploits the dependencies between these tasks through a natural coupling of their interdependencies during inference. In particular, the similarity constraints are relaxed within the tumor regions using an efficient and relatively simple formulation. We evaluated the performance of our formulation both quantitatively and qualitatively for registration and segmentation problems on two publicly available datasets (BraTS 2018 and OASIS 3), reporting competitive results with other recent state-of-the-art methods. Moreover, our proposed framework reports significant amelioration (p < 0.005) for the registration performance inside the tumor locations, providing a generic method that does not need any predefined conditions (e.g., absence of abnormalities) about the volumes to be registered.

code

• https://github.com/TheoEst/joint_registration_tumor_segmentation

---

介绍

脑肿瘤，尤其是胶质瘤，是最常见的类型之一，是最危险和增长最快的癌症类型。放射治疗是治疗脑肿瘤的主要手段之一。然而，放射治疗依赖于医生手动分割肿瘤，使得过程繁琐、耗时，并且由于观察者之间的一致性较低而对偏差敏感。

为了克服这些局限性，最近提出了许多方法，试图提供工具和算法，使胶质瘤分割过程自动化，提升准确率。近年来，大多数利用 BraTS 数据集的深度学习方法都是基于与 VNet 类似的三维卷积神经网络（CNNs）结构。

近年来，人们提出了多种基于深度学习的配准方法，在保证配准精度和鲁棒性的同时，大大减少了计算时间。Dalca¹ 等人提出了一种基于 atlas 的脑 MR 图像配准的深度学习框架 VoxelMorph。然而，当涉及到包含异常的解剖结构时，例如肿瘤区域，由于它们之间缺乏相似性，这些方法无法记录某些位置的体素。这通常会导致变形图像中肿瘤区域及其周围的变形。

为了克服这一问题，本文提出了一种基于深度学习的双任务结构，该结构同时处理配准和肿瘤分割问题，在分割出来的肿瘤区域内放松配准约束，同时提供位移场和分割图。

据我们所知，这是第一次提出基于联合深度学习的体系结构，在两个公开的脑MRI数据集中显示了非常振奋人心的结果。（To the best of our knowledge this is the first time that a joint deep learning-based architecture is presented, showing very promising results in two publicly available datasets for brain MRI.）我们提出的框架通过引入配准分支来阐明大脑解剖学或生理学的临床或功能趋势，提供了一个非常强大的公式（损失函数）。由于与分割分支的协同作用，它进一步实现了脑肿瘤区域的建模和检测。

方法

假设有一对来自两个不同患者的医学 3D 图像——一个浮动图像 S S S，以及一个参考图像 R R R，连同它们的肿瘤区域的分割标签（ S s e g 和 R s e g S_{seg} 和 R_{seg} Sseg​和Rseg​）。我们的框架由一个具有共享参数的双头结构组成，如下图所示。在训练过程中，网络使用上面说的两个图像作为输入，并输出它们的脑肿瘤分割 mask S s e g S_{seg} Sseg​ 和 R s e g R_{seg} Rseg​ 以及将 S S S 投影或映射到 R R R 的最佳可变性变换 G G G。配准部分的目标是找到将 S S S 变换为 R R R 的最佳变形场 G G G。

共享的编码器

我们提出的框架与其他配准方法的主要区别之一是浮动图像和参考图像的组合方式。因为，我们不是串联两个初始图像来一次性输入，而是两次调用同一个编码器独立地输入，分别得到输出浮动图像和参考图像的两组特征映射（称为隐藏编码（latent code）） C s o u r c e C_{source} Csource​ 和 C r e f e r e n c e C_{reference} Creference​。然后， C s o u r c e C_{source} Csource​ 和 C r e f e r e n c e C_{reference} Creference​ 被独立地输入到分割任务解码器中，分别输出得到预测的分割掩膜 S ^ s e g \hat S_{seg} S^seg​ 和 R ^ s e g \hat R_{seg} R^seg​。对于配准任务解码器， C s o u r c e C_{source} Csource​ 和 C r e f e r e n c e C_{reference} Creference​ 会在上图中的 Merge 处进行相减操作，得到的结果输入配准解码器，最后输出从浮动像到参考图像的变形场。

采用这种策略的动机是：我们想要强制编码器从单个图像而不是一对图像中提取分别对配准和分割有意义的特征表示。

• 从配准的角度看

  这相当于要求编码器发现一个不受变形场影响的模板，由此可以忽略配准的变形的影响来对每个图像进行编码（Shu et al.，2018）²，而不是对可能的卷对之间的变形网格进行解码。

• 从分割的角度看

  浮动图像和参考图像的肿瘤图之间没有关系，因此要输入到分割解码器的隐藏编码（latent code）不应该相互依赖。

我们在 Merge 处测试了两种融合运算，即拼接和减法运算。源图像和参考图像都是 4D 体积，其第一维对应于每个受试者使用的 4 种不同的 MRI 模式。从编码器输出得到的 C s o u r c e C_{source} Csource​ 和 C r e f e r e n c e C_{reference} Creference​ 依然是 4D 体积，其中，第一维 n f n_f nf​ 是编码器的最后一块的卷积核的数目。在 C s o u r c e C_{source} Csource​ 和 C r e f e r e n c e C_{reference} Creference​ 输入到配准解码器之前，它们在 Merge 处被合并，其中，

• 在拼接的情况下，输出第一维大小为 2 × n f 2×n_f 2×nf​ 的一个 4D 体积特征图
• 在减法运算情况，下输出第一维大小为 n f n_f nf​ 的一个 4D 体积特征图

两种运算下其余维度都保持不变。特别地，对于每个隐藏编码特征图 C I C_I CI​，减法运算呈现以下自然属性：

• ∀ C I ∈ R n : Merge ⁡ ( C I , C I ) = 0 \forall C_{I} \in \mathbb{R}^{n}: \operatorname{Merge}\left(C_{I}, C_{I}\right)=0 ∀CI​∈Rn:Merge(CI​,CI​)=0
• ∀ C I , C J ∈ R n × R n : Merge ⁡ ( C I , C J ) = − Merge ⁡ ( C J , C I ) \forall C_{I}, C_{J} \in \mathbb{R}^{n} \times \mathbb{R}^{n}: \operatorname{Merge}\left(C_{I}, C_{J}\right)=-\operatorname{Merge}\left(C_{J}, C_{I}\right) ∀CI​,CJ​∈Rn×Rn:Merge(CI​,CJ​)=−Merge(CJ​,CI​)

肿瘤分割解码器

分割解码器会分割出原图像和参考图像的 masks（ S ^ s e g \widehat S_{seg} S seg​ 和 R ^ s e g \widehat R_{seg} R seg​）。但是，这些 masks 指的是无法正确配准的区域，因为没有关于该区域的相应解剖信息。这些信息被整合到配准组件的优化中，用于放松相似性约束（分割出来的肿瘤区域不会用于形变），这在一定程度上保留了肿瘤的几何属性。

在本文中，我们使用加权分类交叉熵损失和 BraTS 数据集提供的针对肿瘤区域的三个不同的分割类别来执行所有的实验。

L s e g = C E ( S s e g , S ^ s e g ) + C E ( R s e g , R ^ s e g ) \mathcal{L}_{seg} = CE(S_{seg}, \widehat S_{seg}) + CE(R_{seg}, \widehat R_{seg}) Lseg​=CE(Sseg​,S seg​)+CE(Rseg​,R seg​)

弹性配准解码器

在本文中，配准策略是基于 Christodoulidis 等人(2018)³ 提出的策略。主要部件是 3D 空间变压器（STN）。STN 利用变形网格 G 对源图像 S 进行变形（或扭曲）。
D = W ( S , G ) D = W(S, G) D=W(S,G)
其中， W ( ⋅ , G ) \mathcal{W}(\cdot, G) W(⋅,G) 表示基于变形网格 G 的采样操作。具体来说，变形网格 G 将馈送到 STN 层作为向后三线圈插值（WTF）采样的采样坐标。

值得一提的是，我们配准解码器输出的是沿每个轴的连续像素之间的空间梯度矩阵 Φ \Phi Φ，而不是直接回归每个像素位移。实际的变形网格 G 可通过沿 x、y 和 z 轴对 Φ \Phi Φ 应用积分运算来获得，该积分运算由离散情况下的累积和来近似。

对于配准，经典的优化方案是最小化 R 和 D 图像强度之间的 Frobenius 范数（Frobenius 范数是针对矩阵而言的，通俗来讲就是矩阵中的元素的平方和再开方）：

L reg  = ∥ ( R − D ) ∥ 2 + α ∥ Φ − Φ I ∥ 1 \mathcal{L}_{\text {reg }}=\|(R-D)\|^{2}+\alpha\left\|\Phi-\Phi_{I}\right\|_{1} Lreg ​=∥(R−D)∥2+α∥Φ−ΦI​∥1​

这里我们要提到的是，在源图像的分割结果 S ^ s e g \widehat S_{seg} S seg​ 上，我们也用了与在源图像 S 上相同的位移网格来使它变形，得到结果 D s e g = W ( S ^ s e g , G ) D_{seg} = \mathcal{W}(\widehat S_{seg}, G) Dseg​=W(S seg​,G)。此外，假设 R ^ s e g 0 \widehat R^0_{seg} R seg0​ ⋅ \cdot ⋅ D s e g 0 D^0_{seg} Dseg0​ 为二元体矩阵，分别表示分割区域之外的体素，以此来决定某区域是否用于配准的训练。然后，配准损失可以写成：

L reg  ⋆ = ∥ ( R − D ) ⋅ D seg  0 ⋅ R ^ seg  0 ∥ 2 + α ∥ Φ − Φ I ∥ 1 \mathcal{L}_{\text {reg }}^{\star}=\left\|(R-D) \cdot D_{\text {seg }}^{0} \cdot \widehat{R}_{\text {seg }}^{0}\right\|^{2}+\alpha\left\|\Phi-\Phi_{I}\right\|_{1} Lreg ⋆​=∥∥∥​(R−D)⋅Dseg 0​⋅R seg 0​∥∥∥​2+α∥Φ−ΦI​∥1​

Φ I \Phi_I ΦI​ 是单位变形的空间梯度，α 是正则化超参数。在位移 Φ \Phi Φ 上使用正则化是必要的，以便约束网络在解剖学上产生更有意义的平滑变形网格，同时正则化可以避免局部最小值。

最后通过分割损失函数和配准损失函数的联合优化，得到最终损失函数：

L = L r e g + β L s e g \mathcal{L} = \mathcal{L}_{reg} + \beta \mathcal{L}_{seg} L=Lreg​+βLseg​

---

我们在这项研究中的贡献有三个方面：

• 联合分割和配准的多任务方案；
• 在编码器之后对隐藏编码（latent code）的融合（merge），以聚集来自该对图像的信息；
• 释放肿瘤区域的配准约束的损失函数 L reg  ⋆ \mathcal{L}_{\text {reg }}^{\star} Lreg ⋆​。

训练细节

对于训练过程，初始学习速率为 2 ⋅ 1 0 − 3 2 \cdot 10^{−3} 2⋅10−3，如果验证集上的性能在 30 个epochs 内没有改善，则初始学习速率降低 5 倍。当 50 个 epochs 没有改善时，训练停止。

正则化权重 α \alpha α 设为 1 0 − 10 10^{-10} 10−10， β \beta β 设为 1。

每 100 个 batches 对网络性能进行评估，观察到两个模型（配准、分割）经过近 200 个 epochs 后收敛。

受限于 GPU 的显存大小，batch_size 设为 2，使用一块 NVIDIA GeForce GTX 1080 Ti GPU 训练大概需要 20 小时。

实验结果

建议自己看原文：https://github.com/TheoEst/joint_registration_tumor_segmentation

总的来说，相对于 Baseline（VNet）的分割性能，这篇论文的方法在单独的分割性能上是没有优势的，因为融合运算（merge）主要用于配准解码器，它会通过损失函数影响编码器的学习，从而间接影响分割解码器，但是这对分割性能的阻碍并不显著。

另一方面，分割提高了配准的精度（voxelmorph 为 0.69，本文为 0.72），毕竟分割区域是肿瘤，而配准解码器的训练忽略了这部分区域的损失和变形，只配准源图像和参考图像的分割区域之外共有的区域。

此外，从我们的实验中我们观察到，肿瘤编码特征的减法在肿瘤区域的配准中表现出更高的性能（被配准变形影响更小）。这表明减法可以为配准任务捕获和编码更多信息性特征。此外，我们对所有基于深度学习的配准方法都实现了很好的泛化，因为我们在一个模型从未见过的数据集（OASIS3 的一部分）中报告了非常稳定的性能。

结论

在这项研究中，我们提出了一个新的基于深度学习的框架来同时处理分割和配准。我们的模型的性能与最近单配准人物任务的最新 SOTA 模型相当（Dalca 等人，2018）¹。然而，我们的模型在肿瘤区域附近的配准有更好的效果。这个效果是通过训练一个共享编码器来实现的，该编码器生成的特征对配准和分割任务都有意义。

同时，这两个问题耦合在一个联合损失函数中，迫使网络聚焦在两个图像中都存在的区域，并同时训练优化这两个问题。

我们在一个联合配准和分割框架中，能够为源图像和参考图像生成高效准确的肿瘤分割图，以及将源图像配准到参考图像空间的位移场。

我们用提出的 L reg  ⋆ \mathcal{L}_{\text {reg }}^{\star} Lreg ⋆​ 做到了保留肿瘤范围的脑部 MRI 配准方法，即，通过这个损失函数，该网络专注于改善共同脑空间中其他组织的局部位移，而不是优化肿瘤区域内的配准损失，根据经验，肿瘤区域是配准误差最大的区域。因此，该网络提高了非肿瘤区域的配准性能。

在未来，我们的目标是使用对抗性损失来评估所提出的框架的性能，以便更好地处理多模态情况。最后，研究自动获取损失函数超参数 α 和 β 的方法。

---

1. Voxelmorph: a learning framework for deformable medical image registratio ↩︎ ↩︎

2. Deforming autoencoders: unsupervised disentangling of shape and appearance ↩︎

3. Christodoulidis, S., Sahasrabudhe, M., et al. (2018). “Linear and deformable image registration with 3D convolutional neural networks, ” in Image Analysis for Moving Organ, Breast, and Thoracic Image, 13-22. ↩︎