【论文笔记】图像分割和图像配准联合学习模型——DeepAtlas

本文是论文《DeepAtlas: Joint Semi-Supervised Learning of Image Registration and Segmentation》的阅读笔记。

文章第一个提出了一个图像配准和图像分割联合学习的网络模型 DeepAtlas，该模型实现了弱监督的图像配准和半监督的图像分割。在图像配准时使用图像的分割标签作为监督数据，如果没有分割标签，则通过分割网络产生；而经过配准后的图像增加了在图像分割时可利用的训练数据的量，相当于是一种数据增强。该模型不仅在分割和配准的精度上有所提升，并且还可以在训练数据有限的情况下实现较好的效果。

一、记号

• $I_m$Im​：浮动图像（moving image）
• $I_t$It​：目标图像（target image）
• $\mathcal{F}_R$FR​：配准网络
• $\theta_r$θr​：配准网络的参数
• $\mathcal{F}_S$FS​：分割网络
• $\theta_s$θs​：分割网络的参数
• $u=\mathcal{F}_R(I_m,I_t;\theta_r)$u=FR​(Im​,It​;θr​)：形变场
• $\phi^{-1}=u+id$ϕ−1=u+id：形变图，其中 $id$id 是恒等变换
• $I_m^w=I_m\circ\phi^{-1}$Imw​=Im​∘ϕ−1：配准后的图像
• $S_t$St​：目标图像分割标签
• $S_m^w=S_m\circ\phi^{-1}$Smw​=Sm​∘ϕ−1：配准后图像分割标签

二、网络结构

DeepAtlas 的目的是当数据集中只有少量的分割标签可用时，通过联合训练来让分割和配准实现较高的精度。

网络的结构如上图所示，蓝色的实线表示弱监督的配准，黄色虚线表示半监督的分割。

文章在附件中给出了分割网络和配准网络的具体结构，如下图左右两图所示：

1. 配准网络

配准网络的损失主要有三个损失函数组成：配准正则损失 $\mathcal{L}_r$Lr​，图像相似度损失 $\mathcal{L}_i$Li​ 和解剖损失（分割相似度损失） $\mathcal{L}_a$La​。配准正则损失 $\mathcal{L}_r$Lr​ 可以让形变场 $\phi$ϕ 变得光滑，图像相似度损失 $\mathcal{L}_i$Li​ 用来评价浮动图像 $I_m$Im​ 和配准后图像 $I_m^w$Imw​ 之间的相似度，解剖损失（分割相似度损失） $\mathcal{L}_a$La​ 是目标图像分割标签 $S_t$St​ 和配准后图像分割标签 $S_m^w$Smw​ 之间的相似度损失。

如此一来，配准学习的过程可以由下式表示：
$\theta_{r}^{\star}=\underset{\theta_{r}}{\operatorname{argmin}}\left\{\mathcal{L}_{i}\left(I_{m} \circ \Phi^{-1}, I_{t}\right)+\lambda_{r} \mathcal{L}_{r}\left(\Phi^{-1}\right)+\lambda_{a} \mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, S_{t}\right)\right\}$θr⋆​=θr​argmin​{Li​(Im​∘Φ−1,It​)+λr​Lr​(Φ−1)+λa​La​(Sm​∘Φ−1,St​)}
其中 $\lambda_r,\lambda_a\geq0$λr​,λa​≥0。

2. 分割网络

分割网络的输入是一张图像 $I$I，输出相应的分割结果 $\hat{S}=\mathcal{F}_S(I;\theta_s)$S^=FS​(I;θs​)，分割网络的损失主要有两个损失函数组成：解剖损失 $\mathcal{L}_a$La​ 和有监督分割损失 $\mathcal{L}_{sp}$Lsp​。解剖损失和配准网络中的相同，有监督的分割损失 $\mathcal{L}_{sp}(\hat{S},S)$Lsp​(S^,S) 是分割网络的分割结果 $\hat{S}$S^ 和人工分割结果 $S$S 之间的相似度损失。但是浮动图像 $I_m$Im​ 和目标图像 $I_t$It​ 的分割标签的存在情况有多种可能，所以相应的损失函数也存在以下四种情况：
$\left\{\begin{array}{l} \mathcal{L}_{a}=\mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, \mathcal{F}_{\mathcal{S}}\left(I_{t}\right)\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right), S_{m}\right), \text { if } I_{t} \text { is unlabeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{a}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right) \circ \Phi^{-1}, S_{t}\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{t}\right), S_{t}\right), \text { if } I_{m} \text { is unlabeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, S_{t}\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right), S_{m}\right), \text { if } I_{m} \text { and } I_{t} \text { are labeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{s p}=0, \text { if both } I_{t} \text { and } I_{m} \text { are unlabeled. } \end{array}\right.$⎩⎪⎪⎨⎪⎪⎧​La​=La​(Sm​∘Φ−1,FS​(It​)) and Lsp​=Lsp​(FS​(Im​),Sm​), if It​ is unlabeled; La​=La​(FS​(Im​)∘Φ−1,St​) and Lsp​=Lsp​(FS​(It​),St​), if Im​ is unlabeled; La​=La​(Sm​∘Φ−1,St​) and Lsp​=Lsp​(FS​(Im​),Sm​), if Im​ and It​ are labeled; La​=Lsp​=0, if both It​ and Im​ are unlabeled. ​
分割的学习过程可以由下式表示：
$\theta_{s}^{\star}=\underset{\theta_{s}}{\operatorname{argmin}}\left(\lambda_{a} \mathcal{L}_{a}+\lambda_{s p} \mathcal{L}_{s p}\right), \quad \lambda_{a}, \lambda_{s p} \geq 0$θs⋆​=θs​argmin​(λa​La​+λsp​Lsp​),λa​,λsp​≥0

三、实施细节

• 解剖相似度损失 $\mathcal{L}_{a}$La​ 和有监督的分割损失 $\mathcal{L}_{sp}$Lsp​ 采用的是 soft multi-class Dice loss：

$\mathcal{L}_{\text {dice}}\left(S, S^{\star}\right)=1-\frac{1}{K} \sum_{k=1}^{K} \frac{\sum_{x} S_{k}(x) S_{k}^{\star}(x)}{\sum_{x} S_{k}(x)+\sum_{x} S_{k}^{\star}(x)}$Ldice​(S,S⋆)=1−K1​k=1∑K​∑x​Sk​(x)+∑x​Sk⋆​(x)∑x​Sk​(x)Sk⋆​(x)​

​ 其中 $k$k 表示分割标签的下标，$x$x 是体素位置，$S$S 和 $S^*$S∗ 是两个要比较的分割标签。

• 图像相似度损失 $\mathcal{L}_i$Li​ 采用的是正则化的互相关（NCC）：

$\mathcal{L}_{i}\left(I_{m}^{w}, I_{t}\right)=1-N C C\left(I_{m}^{w}, I_{t}\right)$Li​(Imw​,It​)=1−NCC(Imw​,It​)

• 配准正则损失 $\mathcal{L}_r$Lr​ 采用的是弯曲能（bending energy）：

$\mathcal{L}_{r}(\mathbf{u})=\frac{1}{N} \sum_{\mathbf{x}} \sum_{i=1}^{d}\left\|H\left(u_{i}(\mathbf{x})\right)\right\|_{F}^{2}$Lr​(u)=N1​x∑​i=1∑d​∥H(ui​(x))∥F2​

​ 其中 $||\cdot||_F$∣∣⋅∣∣F​ 表示弗罗贝尼乌斯范数（Frobenius norm），$H(u_i(x))$H(ui​(x)) 是第 $i$i 个成分 $u(x)$u(x) 的 Hessian 矩阵，$d$d 表示维度，$N$N 表示体素数。

在训练时，会交替的训练分割网络和配准网络，当一个网络在训练时，另一个网络的参数保持不变，并且是每训练配准网络20次才训练分割网络1次，这是因为分割网络更容易收敛。

四、实验结果