本文是文章《Inverse-Consistent Deep Networks for Unsupervised Deformable Image Registration》的阅读笔记。
过去基于学习的配准方法忽略了图像之间转换的逆一致性,并且形变场只被要求局部平滑,不能完全避免形变场的重叠。基于以上两点,提出了一个用于无监督图像配准的逆一致性网络模型ICNet,同时提出了“反重叠约束”来避免形变场的重叠。
一、配准问题
用
S
S
S表示源图像(浮动图像),
T
T
T表示目标图像(固定图像),
F
S
T
F_{ST}
FST表示从图像
S
S
S变到图像
T
T
T的位移场,配准问题可看作是解决以下优化问题:
F
(
S
,
T
)
=
L
similar
(
T
,
Φ
(
S
,
F
S
T
)
)
+
R
(
F
S
T
)
\mathcal{F}(\mathbf{S}, \mathbf{T})=\mathcal{L}_{\text {similar}}\left(\mathbf{T}, \Phi\left(\mathbf{S}, \mathbf{F}_{S T}\right)\right)+\mathcal{R}\left(\mathbf{F}_{S T}\right)
F(S,T)=Lsimilar(T,Φ(S,FST))+R(FST)
其中
Φ
\Phi
Φ表示转换操作,
Φ
(
S
,
F
S
T
)
\Phi(S,F_{ST})
Φ(S,FST)表示根据位移场
F
S
T
F_{ST}
FST变形后的源图像,第一项表示图像之间的相似度,第二项是正则化项。
传统的基于学习的配准算法输出的转换,即位移场或流(flow)通常是不对称的,也就是说图像之间变换的逆一致性被忽略了。这里的逆一致性简单来说就是在配准时不仅要将图像
A
A
A配准到图像
B
B
B,同时还应该将图像
B
B
B配准到图像
A
A
A。
上图中分别是当形变场平滑约束使用较大权重和较小权重时的结果。
二、网络结构
在本文中用
A
A
A表示源图像,
B
B
B表示目标图像,
F
A
B
F_{AB}
FAB表示从图像
A
A
A到图像
B
B
B的位移场,
F
B
A
F_{BA}
FBA表示从图像
B
B
B到图像
A
A
A的位移场。位移场中每个元素是一个三维向量,分别表示在
x
,
y
,
z
x,y,z
x,y,z轴方向的体素从原图像到新图像的位移。变形后的
A
A
A和
B
B
B被记作
A
~
\tilde{A}
A~和
B
~
\tilde{B}
B~。
上图中图(a)是ICNet的网络结构,图(b)是FCN的网络结构,图©是逆网络的结构。
在ICNet中使用两个FCN来预测两个位移场 F A B F_{AB} FAB和 F B A F_{BA} FBA。第一个FCN的目的是根据位移场 F A B F_{AB} FAB将 A A A朝 B B B对齐,并得到变形后的图像 A ~ \tilde{A} A~,第二个FCN的目的是根据位移场 F B A F_{BA} FBA将 B B B朝 A A A对齐,并得到变形后的图像 B ~ \tilde{B} B~。两个FCN的结构相同,参数共享,说白了就是用的同一个FCN。
FCN采用的和UNet类似的结构,n表示FCN中一开始的过滤器的个数,收缩路径的每一步包括一个 3 × 3 × 3 3\times3\times3 3×3×3,步长为1的卷积操作和一个 3 × 3 × 3 3\times3\times3 3×3×3,步长为2的卷积操作(用作下采样)。扩张路径的每一步包括一个 2 × 2 × 2 2\times2\times2 2×2×2,步长为2的反卷积操作用作上采样,并跟着一个拼接操作,将收缩路径和扩张路径的特征图拼接,然后是一个 3 × 3 × 3 3\times3\times3 3×3×3,步长为1的卷积。每个卷积操作后都跟着一个ReLU激活函数,在FCN的最后一层,先用Tanh激活函数将输出值的范围限制在[-1,1],然后再乘以 τ \tau τ,将输出值限制在 [ − τ , τ ] [-\tau, \tau] [−τ,τ], τ \tau τ表示最大的位移。
使用网格采样模块(即空间变换网络,STN)获得变形后的图像,STN采用的是双线性插值,所以是可微的。此外,使用一个逆网络来基于逆一致性损失 L o s s i n v 1 和 L o s s i n v 2 Loss_{inv}1和Loss_{inv}2 Lossinv1和Lossinv2,并根据位移场 F A B F_{AB} FAB生成一个估计的逆位移场 F ~ B A \tilde F_{BA} F~BA。具体的,使用网格采样策略来基于 F A B F_{AB} FAB和 − F A B -F_{AB} −FAB生成逆位移场 F ~ B A \tilde F_{BA} F~BA。对于位移场 F A B F_{AB} FAB,首先获得它的负的位移场 − F A B -F_{AB} −FAB,然后将它们喂入STN中获得估计的逆位移场 F ~ B A \tilde F_{BA} F~BA。同理可获得 F ~ A B \tilde F_{AB} F~AB。
三、损失函数
1. 逆一致性约束
逆一致性约束定义如下:
L
i
n
v
=
∥
F
A
B
−
F
~
A
B
∥
F
2
+
∥
F
B
A
−
F
~
B
A
∥
F
2
F
~
A
B
=
G
(
F
B
A
,
−
F
B
A
)
F
~
B
A
=
G
(
F
A
B
,
−
F
A
B
)
\mathcal{L}_{i n v}=\left\|\mathbf{F}_{A B}-\widetilde{\mathbf{F}}_{A B}\right\|_{F}^{2}+\left\|\mathbf{F}_{B A}-\widetilde{\mathbf{F}}_{B A}\right\|_{F}^{2}\\ \begin{array}{l} \widetilde{\mathbf{F}}_{A B}=\mathcal{G}\left(\mathbf{F}_{B A},-\mathbf{F}_{B A}\right) \\ \widetilde{\mathbf{F}}_{B A}=\mathcal{G}\left(\mathbf{F}_{A B},-\mathbf{F}_{A B}\right) \end{array}
Linv=∥∥∥FAB−F
AB∥∥∥F2+∥∥∥FBA−F
BA∥∥∥F2F
AB=G(FBA,−FBA)F
BA=G(FAB,−FAB)
其中,
G
\mathcal{G}
G是STN产生的映射,
∣
∣
⋅
∣
∣
F
||\cdot||_F
∣∣⋅∣∣F表示矩阵的弗罗贝纽斯(Frobenius)正则。
2. 反重叠约束
反重叠约束:
L
ant
=
Σ
p
∈
Ω
Σ
i
∈
{
x
,
y
,
z
}
δ
(
∇
F
A
B
i
(
p
)
+
1
)
∣
∇
F
A
B
i
(
p
)
∣
2
+
δ
(
∇
F
B
A
i
(
p
)
+
1
)
∣
∇
F
B
A
i
(
p
)
∣
2
\begin{aligned} \left.\mathcal{L}_{\text {ant }}=\Sigma_{p \in \Omega} \Sigma_{i \in\{x, y, z\}}\right. & \delta\left(\nabla \mathbf{F}_{A B}^{i}(p)+1\right)\left|\nabla \mathbf{F}_{A B}^{i}(p)\right|^{2} \\ &+\delta\left(\nabla \mathbf{F}_{B A}^{i}(p)+1\right)\left|\nabla \mathbf{F}_{B A}^{i}(p)\right|^{2} \end{aligned}
Lant =Σp∈ΩΣi∈{x,y,z}δ(∇FABi(p)+1)∣∣∇FABi(p)∣∣2+δ(∇FBAi(p)+1)∣∣∇FBAi(p)∣∣2
其中,
∇
F
A
B
i
(
p
)
\nabla F^i_{AB}(p)
∇FABi(p)是位移场
F
A
B
F_{AB}
FAB中体素
p
p
p在第
i
i
i个轴的梯度,
δ
(
Q
)
\delta(Q)
δ(Q)是用来惩罚位移场中有重叠的位置的梯度的指示函数。如果
Q
≤
0
,
δ
(
Q
)
=
∣
Q
∣
Q\le0,\delta(Q)=|Q|
Q≤0,δ(Q)=∣Q∣,反之,
δ
(
Q
)
=
0
\delta(Q)=0
δ(Q)=0。
如果在位置 p p p沿着第 i i i个轴处有重叠,即 ∇ F A B i ( p ) + 1 ≤ 0 \nabla F^i_{AB}(p)+1\le0 ∇FABi(p)+1≤0,则增加惩罚 ∇ F A B i ( p ) + 1 \nabla F^i_{AB}(p)+1 ∇FABi(p)+1在该位置的梯度上,以让梯度变得小,反之,如果 ∇ F A B i ( p ) + 1 > 0 \nabla F^i_{AB}(p)+1>0 ∇FABi(p)+1>0,即没有重叠,则不做惩罚。
3. 平滑约束
平滑约束的定义如下:
L
s
m
o
=
Σ
p
∈
Ω
(
∥
∇
F
A
B
(
p
)
∥
2
2
+
∥
∇
F
B
A
(
p
)
∥
2
2
)
\mathcal{L}_{s m o}=\Sigma_{p \in \Omega}\left(\left\|\nabla \mathbf{F}_{A B}(p)\right\|_{2}^{2}+\left\|\nabla \mathbf{F}_{B A}(p)\right\|_{2}^{2}\right)
Lsmo=Σp∈Ω(∥∇FAB(p)∥22+∥∇FBA(p)∥22)
其中,
∇
F
A
B
(
p
)
\nabla F_{AB}(p)
∇FAB(p)是体素
p
p
p处的位移场
F
A
B
F_{AB}
FAB的梯度,
∇
F
B
A
(
p
)
\nabla F_{BA}(p)
∇FBA(p)是体素
p
p
p处的位移场
F
B
A
F_{BA}
FBA的梯度,
∣
∣
⋅
∣
∣
2
||\cdot||_2
∣∣⋅∣∣2表示向量的
L
2
L_2
L2正则。
4. 图像相似性度量
使用MSD(mean squared distance)作为图像的相似性度量,该损失用来衡量图像之间的大小差异,其定义如下:
L
s
i
m
=
∥
B
−
A
~
∥
F
2
+
∥
A
−
B
~
∥
F
2
\mathcal{L}_{s i m}=\|\mathbf{B}-\widetilde{\mathbf{A}}\|_{F}^{2}+\|\mathbf{A}-\widetilde{\mathbf{B}}\|_{F}^{2}
Lsim=∥B−A
∥F2+∥A−B
∥F2
其中,
A
~
=
G
(
F
A
B
,
A
)
\widetilde{\mathbf{A}}=\mathcal{G}\left(\mathbf{F}_{A B}, \mathbf{A}\right)
A
=G(FAB,A),
B
~
=
G
(
F
B
A
,
B
)
\widetilde{\mathbf{B}}=\mathcal{G}\left(\mathbf{F}_{B A}, \mathbf{B}\right)
B
=G(FBA,B),
G
\mathcal{G}
G是由STN得到的映射函数。
5. 总损失
ICNet的总损失如下:
L
(
A
,
B
)
=
L
s
i
m
+
α
L
s
m
o
+
β
L
i
n
v
+
γ
L
a
n
t
\begin{aligned} \mathcal{L}(\mathbf{A}, \mathbf{B}) &=\mathcal{L}_{s i m}+\alpha \mathcal{L}_{s m o} \\ &+\beta \mathcal{L}_{i n v}+\gamma \mathcal{L}_{a n t} \end{aligned}
L(A,B)=Lsim+αLsmo+βLinv+γLant
其中
α
,
β
,
γ
\alpha,\beta,\gamma
α,β,γ是平衡因子。
四、实验
1. 预处理
在图像的预处理阶段,要对图像做空间正则化和灰度值正则化。空间正则化包括颅骨去除、小脑移除,并将它们线性对齐到Colin27模板,然后将图像重采样到相同的分辨率( 1 m m × 1 m m × 1 m m 1mm\times1mm\times1mm 1mm×1mm×1mm),再将其剪裁到相同大小( 144 × 192 × 160 144\times192\times160 144×192×160)。灰度值正则化包括先使用直方图匹配算法将灰度直方图匹配到Colin27模板,并进行z-score正则化(减均值,除以标准差),以让图像的均值和标准差分别为0和1。
2. 实验设置
在实现时,使用Adam作为优化器,学习率为 5 e − 4 5e^{-4} 5e−4,使用ADNI-1数据集,随机划分10%的数据作为验证集,剩余的作为训练集,迭代次数为4w次。
在实验中进行两个任务来衡量配准的效果:脑组织分割和解剖学地标检测;实验使用的baseline有Demons、SyN和基于无监督学习的深度学习方法DL(使用MMSE,minimum mean squared error作为相似性度量)。将不适用逆一致性约束的ICNet记为ICNet-1,将不使用反重叠约束的ICNet记为ICNet-2。
在ADNI-2数据集中选用5个MR图像作为图谱,使用多数投票策略进行基于多图谱的分割,以进行脑组织分割任务。在地标检测中,每个图谱中的地标先被映射到指定的测试图像,因此给定测试图像,就得到了5个变形后的地标位置(对于每个地标来说),然后计算5个位置的平均位置得到最终的地标位置。
在ICNet中, γ \gamma γ被设置为 1 0 5 10^5 105, α , β \alpha,\beta α,β被限制在 [ 1 0 − 5 , 1 0 − 4 , . . . , 1 0 5 ] [10^{-5},10^{-4},...,10^5] [10−5,10−4,...,105]范围内。在ICNet-1中 β = 0 , γ = 1 0 5 \beta=0,\gamma=10^5 β=0,γ=105,在ICNet-2中 γ = 0 \gamma=0 γ=0。FCN一开始的过滤器的通道数 n = 8 n=8 n=8, τ = 7 \tau=7 τ=7。在脑组织分割任务中,使用的度量指标为DSC(dice similarity coefficient)、SEN(sensitivity)、PPV(positive predictive value)、ASD(average symmetric surface distance)、HD(Hausdorff distance)。在地标检测任务中,通过计算预测地标位置和ground truth位置之间的欧几里得距离来计算地标检测的误差。ACC、SEN、PPV越高说明效果越好,ASD、HD和检测误差越小说明效果越好。
3. 实验结果
上图是训练和验证的损失,红色的表示训练损失,绿色的表示验证损失,四个图分布是相似度、平滑度、逆一致性和反重叠损失。
上图分别表示原图;三种分割组织图,其中CSF是脑脊液,GM是灰质,WM是白质;五个解剖学地标(用不同颜色的块表示)。
上图是三个脑组织的分割结果,可以发现ICNet在绝大多数分割和衡量指标上都达到了最优结果。
上图是六种算法的配准结果,红色和黄色箭头分别表示左右颞平面。
上图是六个算法在地标检测时的误差,ICNet在3个地标和平均误差中都取得了最优效果。
上图是各个算法的测试时间对比。
上图是不同权重系数下的形变场情况。
上图是两种模型不使用反重叠约束的情况。