本文是论文《CycleMorph: Cycle Consistent Unsupervised Deformable Image Registration》的阅读笔记。
现有的既又深度学习的配准方法通常不能保证拓扑保持(preservation of topology),为了解决这个问题,文章提出了一个名为CycleMorph的模型,使用循环一致性来提供一个隐式的正则来保留拓扑结构。此外,该模型可以实现2D或3D图像以及多尺度图像的配准。
一、相关工作
1. CycleMorph的循环一致性
在该模型中训练两个卷积神经网络 G X G_X GX和 G Y G_Y GY, G X G_X GX产生从源域图像到目标域图像的正向形变场 ϕ X Y \phi_{XY} ϕXY, G Y G_Y GY产生从目标域图像到源域图像的逆向形变场 ϕ Y X \phi_{YX} ϕYX。
为了解决3D图像太大而导致显存放不下的情况,采用了先对下采样的图像进行粗糙的大变形,然后再对局部的patch进行变形。
2. 微分同胚图像配准
在经典的变分图像配准方法中,配准问题的能量方程可以表示为:
L
(
X
,
Y
,
ϕ
)
=
L
s
i
m
(
T
(
X
,
ϕ
)
,
Y
)
+
L
r
e
g
(
ϕ
)
\mathcal{L}(X, Y, \phi)=\mathcal{L}_{s i m}(\mathcal{T}(X, \phi), Y)+\mathcal{L}_{r e g}(\phi)
L(X,Y,ϕ)=Lsim(T(X,ϕ),Y)+Lreg(ϕ)
其中,
X
,
Y
X,Y
X,Y分别表示浮动图像和固定图像,
ϕ
\phi
ϕ表示位移向量场,
T
\mathcal T
T表示根据形变场
ϕ
\phi
ϕ将
X
X
X变形到
Y
Y
Y的变换函数(空间变换层),
L
s
i
m
\mathcal L_{sim}
Lsim是衡量图像相似性的相似项,
L
r
e
g
\mathcal{L}_{reg}
Lreg是使得形变场平滑的正则项。
微分同胚图像配准为向量场 ϕ \phi ϕ施加了一个约束,使得可变形映射是微分同胚。微分同胚变形保证了两个图像之间的变形是连续、可微、拓扑保持的。微分同胚图像配准算法的例子有LDDMM和SyN。
无监督的图像配准方法通常没有为一致性增加一个约束,从而可能会产生形变场重叠的问题。
3. 一致性图像配准
当参数有限时,微分同胚图像配准算法中的形变场通常是离散的。因此,估计的形变 F : X ↦ Y F: X \mapsto Y F:X↦Y和估计的形变 R : Y ↦ X R: Y \mapsto X R:Y↦X的逆是不相等的。在一致性图像配准中,这个问题通过逆一致性 R ≃ F − 1 R\simeq F^{-1} R≃F−1而减轻了。
二、CycleMorph
上图是CycleMorph网络的结构示意图。
1. 损失函数
模型通过优化以下损失函数来训练:
min
G
X
,
G
Y
L
(
X
,
Y
,
G
X
,
G
Y
)
\min _{G_{X}, G_{Y}} \mathcal{L}\left(X, Y, G_{X}, G_{Y}\right)
GX,GYminL(X,Y,GX,GY)
其中
L
(
X
,
Y
,
G
X
,
G
Y
)
=
L
r
e
g
i
s
t
(
X
,
Y
,
G
X
)
+
L
r
e
g
i
s
t
(
Y
,
X
,
G
Y
)
+
α
L
c
y
c
l
e
(
X
,
Y
,
G
X
,
G
Y
)
+
β
L
i
d
e
n
t
i
t
y
(
X
,
Y
,
G
X
,
G
Y
)
\begin{aligned} \mathcal{L}\left(X, Y, G_{X}, G_{Y}\right)=& \mathcal{L}_{r e g i s t}\left(X, Y, G_{X}\right) \\ &+\mathcal{L}_{r e g i s t}\left(Y, X, G_{Y}\right) \\ &+\alpha \mathcal{L}_{c y c l e}\left(X, Y, G_{X}, G_{Y}\right) \\ &+\beta \mathcal{L}_{i d e n t i t y}\left(X, Y, G_{X}, G_{Y}\right) \end{aligned}
L(X,Y,GX,GY)=Lregist(X,Y,GX)+Lregist(Y,X,GY)+αLcycle(X,Y,GX,GY)+βLidentity(X,Y,GX,GY)
上图是网络中三个损失的示意图。
L
r
e
g
i
s
t
\mathcal{L}_{regist}
Lregist是配准损失,它包括图像相似性损失和形变场的正则项损失两部分,图像相似性采用的是局部互相关,正则项损失采用的是
L
2
L_2
L2正则损失,所以
L
regist
(
X
,
Y
,
G
X
)
=
−
(
T
(
X
,
ϕ
X
Y
)
⊗
Y
)
+
λ
∑
∥
∇
ϕ
X
Y
∥
2
\begin{array}{l} \mathcal{L}_{\text {regist}}\left(X, Y, G_{X}\right) \\ \quad=-\left(\mathcal{T}\left(X, \phi_{X Y}\right) \otimes Y\right)+\lambda \sum\left\|\nabla \phi_{X Y}\right\|^{2} \end{array}
Lregist(X,Y,GX)=−(T(X,ϕXY)⊗Y)+λ∑∥∇ϕXY∥2
其中
λ
\lambda
λ是超参数,
⊗
\otimes
⊗表示局部互相关,其定义为:
A
⊗
B
=
∑
v
∈
Ω
(
∑
v
i
(
A
(
v
i
)
−
A
ˉ
(
v
)
)
(
B
(
v
i
)
−
B
ˉ
(
v
)
)
)
2
(
∑
v
i
(
A
(
v
i
)
−
A
ˉ
(
v
)
)
2
)
(
∑
v
i
(
B
(
v
i
)
−
B
ˉ
(
v
)
)
2
)
A \otimes B=\sum_{\mathbf{v} \in \Omega} \frac{\left(\sum_{\mathbf{v}_{i}}\left(A\left(\mathbf{v}_{i}\right)-\bar{A}(\mathbf{v})\right)\left(B\left(\mathbf{v}_{i}\right)-\bar{B}(\mathbf{v})\right)\right)^{2}}{\left(\sum_{\mathbf{v}_{i}}\left(A\left(\mathbf{v}_{i}\right)-\bar{A}(\mathbf{v})\right)^{2}\right)\left(\sum_{\mathbf{v}_{i}}\left(B\left(\mathbf{v}_{i}\right)-\bar{B}(\mathbf{v})\right)^{2}\right)}
A⊗B=v∈Ω∑(∑vi(A(vi)−Aˉ(v))2)(∑vi(B(vi)−Bˉ(v))2)(∑vi(A(vi)−Aˉ(v))(B(vi)−Bˉ(v)))2
其中,
Ω
\Omega
Ω表示整个3D图像,
A
ˉ
(
v
)
\bar{A}(v)
Aˉ(v)和
B
ˉ
(
v
)
\bar{B}(v)
Bˉ(v)表示图像
A
(
v
)
A(v)
A(v)和
B
(
v
)
B(v)
B(v)的局部均值,
v
i
v_i
vi迭代在体素
v
v
v周围大小为
w
×
w
×
w
w\times w\times w
w×w×w(2D图像时为
w
×
w
w\times w
w×w)的范围迭代计算,其中
w
=
9
w=9
w=9。
L
c
y
c
l
e
\mathcal{L}_{cycle}
Lcycle是循环一致性损失,源域图像
X
X
X先通过
G
X
G_X
GX被配准得到
Y
^
\hat{Y}
Y^,再通过
G
Y
G_Y
GY配准得到
X
~
\tilde{X}
X~,计算
X
X
X和
X
~
\tilde{X}
X~之间的循环一致性损失;同理目标域图像
Y
Y
Y先通过
G
Y
G_Y
GY被配准得到
X
^
\hat{X}
X^,再通过
G
X
G_X
GX配准得到
Y
~
\tilde{Y}
Y~,计算
Y
Y
Y和
Y
~
\tilde{Y}
Y~之间的循环一致性损失。上述过程可以表示为:
(
X
,
Y
)
≃
(
T
(
Y
^
,
ϕ
^
Y
X
)
,
T
(
X
^
,
ϕ
^
X
Y
)
)
(X, Y) \simeq\left(\mathcal{T}\left(\hat{Y}, \hat{\phi}_{Y X}\right), \mathcal{T}\left(\hat{X}, \hat{\phi}_{X Y}\right)\right)
(X,Y)≃(T(Y^,ϕ^YX),T(X^,ϕ^XY))
其中
(
Y
^
,
X
^
)
:
=
(
T
(
X
,
ϕ
X
Y
)
,
T
(
Y
,
ϕ
Y
X
)
)
(\hat{Y}, \hat{X}):=\left(\mathcal{T}\left(X, \phi_{X Y}\right), \mathcal{T}\left(Y, \phi_{Y X}\right)\right)
(Y^,X^):=(T(X,ϕXY),T(Y,ϕYX))
所以,循环一致性损失可以通过下式计算:
L
cycle
(
X
,
Y
,
G
X
,
G
Y
)
=
∥
T
(
Y
^
,
ϕ
^
Y
X
)
−
X
∥
1
+
∥
T
(
X
^
,
ϕ
^
X
Y
)
−
Y
∥
1
\begin{array}{l} \mathcal{L}_{\text {cycle}}\left(X, Y, G_{X}, G_{Y}\right) \\ \quad=\left\|\mathcal{T}\left(\hat{Y}, \hat{\phi}_{Y X}\right)-X\right\|_{1}+\left\|\mathcal{T}\left(\hat{X}, \hat{\phi}_{X Y}\right)-Y\right\|_{1} \end{array}
Lcycle(X,Y,GX,GY)=∥∥∥T(Y^,ϕ^YX)−X∥∥∥1+∥∥∥T(X^,ϕ^XY)−Y∥∥∥1
其中,
∥
⋅
∥
1
\|\cdot\|_1
∥⋅∥1表示
L
1
L_1
L1范式。
L
i
d
e
n
t
i
t
y
\mathcal{L}_{identity}
Lidentity是恒等损失表示当固定图像和浮动图像为同一个时,两者之间不会有任何变形。因此该损失避免了不必要的变形,使得形变场在静态区域的预测更加稳定。该损失可以表示为:
L
identity
(
X
,
Y
,
G
X
,
G
Y
)
=
−
(
T
(
Y
,
G
X
(
Y
,
Y
)
)
⊗
Y
)
−
(
T
(
X
,
G
Y
(
X
,
X
)
)
⊗
X
)
\begin{array}{l} \mathcal{L}_{\text {identity}}\left(X, Y, G_{X}, G_{Y}\right) \\ =-\left(\mathcal{T}\left(Y, G_{X}(Y, Y)\right) \otimes Y\right)-\left(\mathcal{T}\left(X, G_{Y}(X, X)\right) \otimes X\right) \end{array}
Lidentity(X,Y,GX,GY)=−(T(Y,GX(Y,Y))⊗Y)−(T(X,GY(X,X))⊗X)
其中
⊗
\otimes
⊗表示局部互相关。
2. 空间变换层
对于3D图像来说,使用三线性插值,可以表示为:
T
(
X
,
ϕ
)
=
∑
q
∈
N
(
p
+
ϕ
(
p
)
)
X
(
q
)
∏
d
∈
{
i
,
j
,
k
}
(
1
−
∣
p
d
+
ϕ
(
p
d
)
−
p
d
∣
)
\begin{array}{l} \mathcal{T}(X, \phi) \\ =\sum_{q \in \mathcal{N}(p+\phi(p))} X(q) \prod_{d \in\{i, j, k\}}\left(1-\left|p_{d}+\phi\left(p_{d}\right)-p_{d}\right|\right) \end{array}
T(X,ϕ)=∑q∈N(p+ϕ(p))X(q)∏d∈{i,j,k}(1−∣pd+ϕ(pd)−pd∣)
其中,
p
p
p表示像素下标,
N
(
p
+
ϕ
(
p
)
)
\mathcal{N}(p+\phi(p))
N(p+ϕ(p))表示在
p
+
ϕ
(
p
)
p+\phi(p)
p+ϕ(p)周围的8像素的立方体邻居,
d
d
d是3D图像空间的三个方向。当为2D图像配准时,采用的是双线性插值。
3. 多尺度配准
上图是多尺度配准的示意图,训练时多尺度配准采用先做全局配准,后做局部配准。全局图像配准模型首先在下采样的图像对上进行训练,然后通过对形变场进行上采样得到全尺寸的变形后的图像,然后对变形后图像的patch进行局部图像配准的训练。
在测试阶段,由于要进行两次线性插值,所以两个配准网络会导致误差的累计,从而降低配准的精度。所以采取了全局配准和局部配准先得到各自的形变场,然后再将形变场相加,只做一次线性插值。
三、实验
1. 数据集
三个任务三个数据集,2D配准任务使用的是面部表情图像Radboud Faces Database(RaFD)数据集,3D配准任务使用的是OASIS-3脑部MRI数据集,多尺度配准使用的是肝脏CT数据集。不同任务所采用的参数就不再赘述。
实验采用的baseline有Elastix,SyN,VoxelMorph和MS-DIRNet。
2. 评价指标
使用形变场的雅克比行列式中非负数的百分比来衡量形变场的重叠程度:
∣
J
ϕ
(
v
)
∣
=
∣
∇
ϕ
(
v
)
∣
≤
0
\left|J_{\phi}(\mathbf{v})\right|=|\nabla \phi(\mathbf{v})| \leq 0
∣Jϕ(v)∣=∣∇ϕ(v)∣≤0
其中
v
v
v表示体素位置,
∣
⋅
∣
|\cdot|
∣⋅∣表示矩阵的行列式,当雅克比矩阵的行列式的值都是正数时,形变场是微分同胚的。
在面部表情数据集中有ground-truth标签,衡量其配准后的归一化均方误差损失(NMSE)和结构相似性损失(SSIM);在脑部MR数据集中有脑部结构的分割标签,衡量其配准后的DICE值,其表达式如下:
Dice
(
A
,
B
)
=
2
T
P
2
T
P
+
F
P
+
F
N
\operatorname{Dice}(A, B)=\frac{2 T P}{2 T P+F P+F N}
Dice(A,B)=2TP+FP+FN2TP
其中
T
P
,
F
P
,
F
N
TP,FP,FN
TP,FP,FN分别表示正阳例、负阳例、负阴例像素的个数。
在肝脏CT数据集中有地标点信息,衡量其在配准后的目标配准损失(TRE),TRE可以通过平均欧几里得距离来计算:
T
R
E
(
A
,
B
)
=
1
N
∑
i
=
1
N
∥
a
i
−
b
i
∥
T R E(A, B)=\frac{1}{N} \sum_{i=1}^{N}\left\|a_{i}-b_{i}\right\|
TRE(A,B)=N1i=1∑N∥ai−bi∥
其中
N
N
N是地标点的数目,
a
i
,
b
i
a_i,b_i
ai,bi是浮动图像和固定图像中第
i
i
i个地标坐标向量。
3. 实验结果
上图是在面部表情数据集上的实验结果。
上图是在脑部MR数据上的配准结果。
上图是不同配准方法在脑部MR数据集上的配准结果的对比表。