本文是关于《DEFORM-GAN:AN UNSUPERVISED LEARNING MODEL FOR DEFORMABLE REGISTRATION》论文的阅读笔记。
一、简介
文章提出了一种基于 GAN 的三维医学图像跨模态配准模型 Deform-GAN,并首次将梯度损失(gradient loss)应用到基于深度学习的配准模型中。模型是无监督的,不需要任何 ground-truth 或人工标注信息。
基于无监督的学习模型存在着两个挑战,一个是跨模态或跨序列配准中损失函数的确定,另一个是没有 ground-truth,这使得 GAN 模型比较难训练。
上图是 Deform-GAN 和其他模型的配准结果对比图。
二、记号
- T T T:图像变换网络
- G G G:图像生成器
- D D D:判别器
- R R R:参考图像
- F F F:浮动图像
- ϕ \phi ϕ:配准形变场
- F ( ϕ ) F(\phi) F(ϕ):变形后的浮动图像
- p p p:体素位置
- I I I:三维图像
三、网络结构和训练
上图是 Deform-GAN 网络的结构示意图。
网络主要由三部分组成:图像变换网络 T T T,图像生成器 G G G 和判别器 D D D,前两者采用的都是类似于 U-Net 的网络结构。图像变换网络 T T T 的输入是参考图像和浮动图像的图像对,输出是配准形变场 ϕ \phi ϕ, ϕ \phi ϕ 会进一步通过空间变换器对浮动图像 T T T 做变形得到 F ( ϕ ) F(\phi) F(ϕ),这一步相当于实现了模态间配准(从源域到源域)。图像生成器 G G G 以 F ( ϕ ) F(\phi) F(ϕ) 为输入,其输出是一个接近于参考图像 R R R 的图像 F ′ ( ϕ ) F'(\phi) F′(ϕ),这一步相当于实现了跨模态配准(从源域到目标域)。判别器 D D D 用来判别图像是真实图像还是合成图像。如此一来,配准问题就被分成了两部分:多模配准和单模配准。
在训练前期, T T T 还没有学习到很好的特征,配准效果不好,如果直接将 R R R 和 F ′ ( ϕ ) F'(\phi) F′(ϕ) (这里论文中写的是 F ( ϕ ) F(\phi) F(ϕ),我怀疑是写错了,按照示网络结构示意图中的来)输入到判别器则会产生错误对齐的 F ′ ( ϕ ) F'(\phi) F′(ϕ)。为解决该问题,文章提出了“梯度约束的 GAN 方法”,这种方法的特点是损失函数不是固定的,而是通过学习获得的,它不仅惩罚输出图像和目标域图像之间的差别,而惩罚输出图像和源域图像之间的差别。
生成器的任务有三个:一是骗过判别器,二是最小化输出图像和目标域图像的 L 1 L_1 L1 距离,三是保持输出图像和源域图像在细节上保持相似。
在训练时,三个网络的训练顺序是先训练判别器 D D D,再训练生成器 G G G,最后训练图像变换网络 T T T,当训练一个网络时,其他网络的参数保持不变。
四、损失函数
局部梯度计算公式:
∇
I
^
(
p
)
=
(
∑
p
∈
n
3
x
′
(
p
)
,
∑
p
∈
n
3
y
′
(
p
)
,
∑
p
∈
n
3
z
′
(
p
)
)
\nabla \hat{I}(p)=\left(\sum_{p \in n^{3}} x^{\prime}(p), \sum_{p \in n^{3}} y^{\prime}(p), \sum_{p \in n^{3}} z^{\prime}(p)\right)
∇I^(p)=⎝⎛p∈n3∑x′(p),p∈n3∑y′(p),p∈n3∑z′(p)⎠⎞
n
n
n 是
p
p
p 周围的像素点个数,当
n
n
n 较小时网络难收敛,当较大时图像
R
R
R 和
F
F
F 的边缘很难准确对齐,通过尝试当
n
=
7
n=7
n=7 时效果最好。根据以上梯度的计算公式就可以得到正则化的梯度了:
n
(
I
,
p
)
=
∇
I
^
(
p
)
∥
∇
I
^
(
p
)
∥
+
ε
n(I, p)=\frac{\nabla \hat{I}(p)}{\|\nabla \hat{I}(p)\|+\varepsilon}
n(I,p)=∥∇I^(p)∥+ε∇I^(p)
其中
∣
∣
⋅
∣
∣
||\cdot||
∣∣⋅∣∣ 表示
L
2
L_2
L2 距离,图像
R
R
R 和
F
F
F 之间的局部梯度损失就可以定义为:
L
L
G
(
R
,
F
)
=
∑
p
∈
Ω
∣
n
(
R
,
p
)
⋅
n
(
F
,
p
)
∣
L_{L G}(R, F)=\sum_{p \in \Omega}|n(R, p) \cdot n(F, p)|
LLG(R,F)=p∈Ω∑∣n(R,p)⋅n(F,p)∣
其中
Ω
\Omega
Ω 是
R
R
R 和
F
F
F 的图像域,
图像转换网络
T
T
T 的损失可以表示为:
L
T
(
R
,
F
,
ϕ
)
=
L
s
i
m
(
R
,
F
(
ϕ
)
)
+
α
L
s
m
o
o
t
h
(
ϕ
)
L_{T}(R, F, \phi)=L_{s i m}(R, F(\phi))+\alpha L_{s m o o t h}(\phi)
LT(R,F,ϕ)=Lsim(R,F(ϕ))+αLsmooth(ϕ)
其中
L
s
i
m
(
R
,
F
(
ϕ
)
)
L_{sim}(R,F(\phi))
Lsim(R,F(ϕ)) 又由两部分组成:图像
R
R
R 和
F
′
(
ϕ
)
F'(\phi)
F′(ϕ) 之间负的局部互信息和图像
R
R
R 和
F
(
ϕ
)
F(\phi)
F(ϕ) 之间负的局部梯度距离,如下式所示:
L
s
i
m
(
R
,
F
(
ϕ
)
)
=
−
L
L
C
C
(
R
,
F
′
(
ϕ
)
)
−
β
L
L
G
(
R
,
F
(
ϕ
)
)
L_{s i m}(R, F(\phi))=-L_{L C C}\left(R, F^{\prime}(\phi)\right)-\beta L_{L G}(R, F(\phi))
Lsim(R,F(ϕ))=−LLCC(R,F′(ϕ))−βLLG(R,F(ϕ))
形变场的平滑损失
L
s
m
o
o
t
h
(
ϕ
)
L_{smooth}(\phi)
Lsmooth(ϕ) 可以表示为:
L
smooth
(
ϕ
)
=
∑
p
∈
Ω
∥
∇
ϕ
(
p
)
∥
2
L_{\text {smooth}}(\phi)=\sum_{p \in \Omega}\|\nabla \phi(p)\|^{2}
Lsmooth(ϕ)=p∈Ω∑∥∇ϕ(p)∥2
下面来看下生成器
G
G
G 和判别器
D
D
D 的损失函数。在 Pix2Pix 网络,它的损失函数为:
L
G
∗
=
arg
min
G
max
D
L
c
G
A
N
(
G
,
D
)
+
λ
L
L
1
(
G
)
L_{G^{*}}=\arg \min _{G} \max _{D} L_{c^{G A N}}(G, D)+\lambda L_{L 1}(G)
LG∗=argGminDmaxLcGAN(G,D)+λLL1(G)
其中
L
c
G
A
N
L_{cGAN}
LcGAN 是条件 GAN(conditional GAN)的目标函数,
L
L
1
L_{L1}
LL1 是源图像和 ground-truth 的目标图像之间的
L
1
L_1
L1 距离。但是在本论文中由于源图像和目标图像不是像素级的映射数据,所以不适用。本文用局部梯度损失来限制合成图像
F
′
(
ϕ
)
F'(\phi)
F′(ϕ) 和源图像
F
(
ϕ
)
F(\phi)
F(ϕ) 之间的梯度距离,并保证最终的输出图像的细节和源图像一致。所以最终 GAN 的总损失为:
L
G
′
=
arg
min
G
max
D
L
c
G
A
N
(
G
,
D
)
−
μ
L
L
G
(
F
′
(
ϕ
)
,
F
(
ϕ
)
)
+
λ
L
L
1
(
F
′
(
ϕ
)
,
R
)
\begin{aligned} L_{G^{\prime}}=& \arg \min _{G} \max _{D} L_{c^{G A N}}(G, D)-\mu L_{L G}\left(F^{\prime}(\phi), F(\phi)\right) \\ &+\lambda L_{L 1}\left(F^{\prime}(\phi), R\right) \end{aligned}
LG′=argGminDmaxLcGAN(G,D)−μLLG(F′(ϕ),F(ϕ))+λLL1(F′(ϕ),R)