本文是论文《DeepAtlas: Joint Semi-Supervised Learning of Image Registration and Segmentation》的阅读笔记。
文章第一个提出了一个图像配准和图像分割联合学习的网络模型 DeepAtlas,该模型实现了弱监督的图像配准和半监督的图像分割。在图像配准时使用图像的分割标签作为监督数据,如果没有分割标签,则通过分割网络产生;而经过配准后的图像增加了在图像分割时可利用的训练数据的量,相当于是一种数据增强。该模型不仅在分割和配准的精度上有所提升,并且还可以在训练数据有限的情况下实现较好的效果。
一、记号
- I m I_m Im:浮动图像(moving image)
- I t I_t It:目标图像(target image)
- F R \mathcal{F}_R FR:配准网络
- θ r \theta_r θr:配准网络的参数
- F S \mathcal{F}_S FS:分割网络
- θ s \theta_s θs:分割网络的参数
- u = F R ( I m , I t ; θ r ) u=\mathcal{F}_R(I_m,I_t;\theta_r) u=FR(Im,It;θr):形变场
- ϕ − 1 = u + i d \phi^{-1}=u+id ϕ−1=u+id:形变图,其中 i d id id 是恒等变换
- I m w = I m ∘ ϕ − 1 I_m^w=I_m\circ\phi^{-1} Imw=Im∘ϕ−1:配准后的图像
- S t S_t St:目标图像分割标签
- S m w = S m ∘ ϕ − 1 S_m^w=S_m\circ\phi^{-1} Smw=Sm∘ϕ−1:配准后图像分割标签
二、网络结构
DeepAtlas 的目的是当数据集中只有少量的分割标签可用时,通过联合训练来让分割和配准实现较高的精度。
网络的结构如上图所示,蓝色的实线表示弱监督的配准,黄色虚线表示半监督的分割。
文章在附件中给出了分割网络和配准网络的具体结构,如下图左右两图所示:
1. 配准网络
配准网络的损失主要有三个损失函数组成:配准正则损失 L r \mathcal{L}_r Lr,图像相似度损失 L i \mathcal{L}_i Li 和解剖损失(分割相似度损失) L a \mathcal{L}_a La。配准正则损失 L r \mathcal{L}_r Lr 可以让形变场 ϕ \phi ϕ 变得光滑,图像相似度损失 L i \mathcal{L}_i Li 用来评价浮动图像 I m I_m Im 和配准后图像 I m w I_m^w Imw 之间的相似度,解剖损失(分割相似度损失) L a \mathcal{L}_a La 是目标图像分割标签 S t S_t St 和配准后图像分割标签 S m w S_m^w Smw 之间的相似度损失。
如此一来,配准学习的过程可以由下式表示:
θ
r
⋆
=
argmin
θ
r
{
L
i
(
I
m
∘
Φ
−
1
,
I
t
)
+
λ
r
L
r
(
Φ
−
1
)
+
λ
a
L
a
(
S
m
∘
Φ
−
1
,
S
t
)
}
\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⋆=θrargmin{Li(Im∘Φ−1,It)+λrLr(Φ−1)+λaLa(Sm∘Φ−1,St)}
其中
λ
r
,
λ
a
≥
0
\lambda_r,\lambda_a\geq0
λr,λa≥0。
2. 分割网络
分割网络的输入是一张图像
I
I
I,输出相应的分割结果
S
^
=
F
S
(
I
;
θ
s
)
\hat{S}=\mathcal{F}_S(I;\theta_s)
S^=FS(I;θs),分割网络的损失主要有两个损失函数组成:解剖损失
L
a
\mathcal{L}_a
La 和有监督分割损失
L
s
p
\mathcal{L}_{sp}
Lsp。解剖损失和配准网络中的相同,有监督的分割损失
L
s
p
(
S
^
,
S
)
\mathcal{L}_{sp}(\hat{S},S)
Lsp(S^,S) 是分割网络的分割结果
S
^
\hat{S}
S^ 和人工分割结果
S
S
S 之间的相似度损失。但是浮动图像
I
m
I_m
Im 和目标图像
I
t
I_t
It 的分割标签的存在情况有多种可能,所以相应的损失函数也存在以下四种情况:
{
L
a
=
L
a
(
S
m
∘
Φ
−
1
,
F
S
(
I
t
)
)
and
L
s
p
=
L
s
p
(
F
S
(
I
m
)
,
S
m
)
,
if
I
t
is unlabeled;
L
a
=
L
a
(
F
S
(
I
m
)
∘
Φ
−
1
,
S
t
)
and
L
s
p
=
L
s
p
(
F
S
(
I
t
)
,
S
t
)
,
if
I
m
is unlabeled;
L
a
=
L
a
(
S
m
∘
Φ
−
1
,
S
t
)
and
L
s
p
=
L
s
p
(
F
S
(
I
m
)
,
S
m
)
,
if
I
m
and
I
t
are labeled;
L
a
=
L
s
p
=
0
,
if both
I
t
and
I
m
are unlabeled.
\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.
分割的学习过程可以由下式表示:
θ
s
⋆
=
argmin
θ
s
(
λ
a
L
a
+
λ
s
p
L
s
p
)
,
λ
a
,
λ
s
p
≥
0
\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⋆=θsargmin(λaLa+λspLsp),λa,λsp≥0
三、实施细节
- 解剖相似度损失 L a \mathcal{L}_{a} La 和有监督的分割损失 L s p \mathcal{L}_{sp} Lsp 采用的是 soft multi-class Dice loss:
L dice ( S , S ⋆ ) = 1 − 1 K ∑ k = 1 K ∑ x S k ( x ) S k ⋆ ( x ) ∑ x S k ( x ) + ∑ x S k ⋆ ( x ) \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−K1k=1∑K∑xSk(x)+∑xSk⋆(x)∑xSk(x)Sk⋆(x)
其中 k k k 表示分割标签的下标, x x x 是体素位置, S S S 和 S ∗ S^* S∗ 是两个要比较的分割标签。
- 图像相似度损失 L i \mathcal{L}_i Li 采用的是正则化的互相关(NCC):
L i ( I m w , I t ) = 1 − N C C ( I m w , I t ) \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)
- 配准正则损失 L r \mathcal{L}_r Lr 采用的是弯曲能(bending energy):
L r ( u ) = 1 N ∑ x ∑ i = 1 d ∥ H ( u i ( x ) ) ∥ F 2 \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)=N1x∑i=1∑d∥H(ui(x))∥F2
其中 ∣ ∣ ⋅ ∣ ∣ F ||\cdot||_F ∣∣⋅∣∣F 表示弗罗贝尼乌斯范数(Frobenius norm), H ( u i ( x ) ) H(u_i(x)) H(ui(x)) 是第 i i i 个成分 u ( x ) u(x) u(x) 的 Hessian 矩阵, d d d 表示维度, N N N 表示体素数。
在训练时,会交替的训练分割网络和配准网络,当一个网络在训练时,另一个网络的参数保持不变,并且是每训练配准网络20次才训练分割网络1次,这是因为分割网络更容易收敛。