本文是论文《Cycle-Consistent Training for Reducing Negative Jacobian Determinant in Deep Registration Networks》的阅读笔记。
微分同胚
配准任务中使用的变换(形变场)理想情况下应该是微分同胚的,这样图像的拓扑性在转换的时候就不会发生变化。微分同胚即是可逆和光滑的,现有的配准方法通常不能从理论上保证形变场的平滑性和可逆性。
之前的无监督的深度学习方法是通过使用局部平滑约束来直接使形变场保持平滑,或者对网络结构进行设计来保证形变场平滑。文章通过研究可能的训练机制/任务,这些机制/任务将帮助网络避免使用负雅克比行列式预测变换并产生更平滑的变形。
上图是配准的大体框架(VoxelMorph就是采用以上结构),其中采样器,即空间转换网络(STN)由变形单元(deformation unit)和采样单元(sampling unit)组成,假设源域图像和目标域图像分别为 x , y x,y x,y,则变形单元会产生一个静态位移场 u u u,那么就可以根据位移场 u u u,使用插值的方法从源图像 x x x得到采样后的图像 y ~ = x ( I d + u ) \tilde{y}=x(Id+u) y~=x(Id+u),其中 I d Id Id表示恒等映射。
微分同胚
ϕ
\phi
ϕ对图像
x
x
x的操作被近似为
ϕ
⋅
x
=
x
∘
ϕ
−
1
≈
x
(
I
d
+
u
)
\phi \cdot x=x \circ \phi^{-1} \approx x(I d+\mathbf{u})
ϕ⋅x=x∘ϕ−1≈x(Id+u),位移场
u
u
u的平滑性通常通过对其导数
D
u
Du
Du添加正则化来保证。比如VoxelMorph网络的损失函数为:
C
C
(
y
,
y
~
)
+
λ
∥
D
u
∥
l
2
C C(y, \tilde{y})+\lambda\|D \mathbf{u}\|_{l_{2}}
CC(y,y~)+λ∥Du∥l2,其中
C
C
CC
CC表示互相关,
λ
\lambda
λ是控制正则项的权重。
形变场的雅克比行列式,即 det ( D ϕ − 1 ) ≈ det ( I d + D u ) \operatorname{det}\left(D \phi^{-1}\right) \approx \operatorname{det}(\mathrm{Id}+D \mathbf{u}) det(Dϕ−1)≈det(Id+Du)很有可能在多个位置都是负的,即使增大 λ \lambda λ的值时,这种重叠问题依然可能会存在。此外,超参 λ \lambda λ的确定也是比较麻烦的,比较大的 λ \lambda λ值会导致较小的位移场,而降低模型的精度。如上图所示,分别是 λ = 1 , 2 , 4 \lambda=1,2,4 λ=1,2,4时形变场的折叠情况。
Prob-VoxelMorph采用一个重参数的技巧,并插入集成层(integration layer)来让模型产生平滑的形变场。循环一致性是通过为模型增加一个额外的任务来间接的解决形变场的平滑性问题。
循环一致性
不同于以往的模型,文本不是通过增加一个特定的损失项来防止形变场重叠,而是通过研究有可能的训练机制/任务来实现。循环一致性的想法不会改变backbone网络的结构,而是通过增加一个从已经预测的图像中恢复原始图像的任务来实现正则化。
循环一致性具体来说就是,对于图像
x
,
y
x,y
x,y,先将
y
y
y作为固定图像,将
x
x
x配准到
y
y
y,得到配准图像
y
~
\tilde y
y~,该变换的位移场为
u
x
→
y
~
u_{x\rightarrow\tilde y}
ux→y~;然后再以
x
x
x作为固定图像,将
y
~
\tilde y
y~配准到图像
x
x
x,该变换的位移场为
u
y
~
→
x
~
u_{\tilde y\rightarrow\tilde x}
uy~→x~。上图是循环一致性配准的流程示意图。整个网络的循环一致性损失为:
C
C
(
y
,
y
~
)
+
λ
∥
D
u
x
→
y
~
∥
l
2
+
C
C
(
x
,
x
~
)
+
λ
∥
D
u
y
~
→
x
~
∥
l
2
C C(y, \tilde{y})+\lambda\left\|D \mathbf{u}_{x \rightarrow \tilde{y}}\right\|_{l_{2}}+C C(x, \tilde{x})+\lambda\left\|D \mathbf{u}_{\tilde{y} \rightarrow \tilde{x}}\right\|_{l_{2}}
CC(y,y~)+λ∥Dux→y~∥l2+CC(x,x~)+λ∥Duy~→x~∥l2
这种设计直接解决了网络的可逆性,循环一致性任务同时也间接保证了网络可以学习到一个平滑的形变场,即帮助网络排除那些非循环一致的变换。此外这种设计并不会增加模型的参数。
循环一致性的想法虽然跟双向配准有些相似,但是还是有很大的不同。在循环一致性配准中,目标图像从没发生变形。双向配准是既使用 ( x , y ) (x,y) (x,y)又使用 ( y , x ) (y,x) (y,x)作为输入,而循环一致性配准只使用 ( x , y ) (x,y) (x,y)作为输入。上表是两种方式的对比。
实验
实验使用的是MindBoggle101数据集,该数据集包括三个子数据集:
- NKI-RS-22
- NKI-TRT-20
- OASIS-TRT-20
每个图像都是
182
×
218
×
182
182\times218\times182
182×218×182大小的,然后裁剪到
144
×
180
×
144
144\times180\times144
144×180×144大小,这些图像已经对齐到MNI152空间,还需要将其灰度值归一化到[0, 1]的范围。原始的VoxelMorph将作为本文模型的backbone网络,实验具体设置为
λ
=
1
\lambda=1
λ=1,epoch=10,选用学习率为
1
e
−
4
1e^{-4}
1e−4的Adam优化器,使用三折交叉验证。使用Dice值作为评价指标,其表达式如下:
Dice
(
(
ϕ
⋅
L
x
i
)
,
L
y
i
)
=
2
∣
(
ϕ
⋅
L
x
i
)
∩
L
y
i
∣
∣
ϕ
⋅
L
x
i
∣
+
∣
L
y
i
∣
\operatorname{Dice}\left(\left(\phi \cdot L_{x}^{i}\right), L_{y}^{i}\right)=\frac{2\left|\left(\phi \cdot L_{x}^{i}\right) \cap L_{y}^{i}\right|}{\left|\phi \cdot L_{x}^{i}\right|+\left|L_{y}^{i}\right|}
Dice((ϕ⋅Lxi),Lyi)=∣ϕ⋅Lxi∣+∣∣Lyi∣∣2∣∣(ϕ⋅Lxi)∩Lyi∣∣
其中
L
x
i
L^i_x
Lxi是带有下标的标签,Dice值越大说明配准效果越好。
使用OASIS-TRT-20作为测试集。形变场的重叠是通过检查负雅可比行列式发生的位置来获得的。设
P
P
P定义为所有体素
V
V
V上Jacobian行列式为负的体素位置的百分比,即:
P
:
=
∑
δ
(
det
(
D
ϕ
−
1
)
<
0
)
V
\mathcal{P}:=\frac{\sum \delta\left(\operatorname{det}\left(D \phi^{-1}\right)<0\right)}{V}
P:=V∑δ(det(Dϕ−1)<0)
理想情况下,
P
P
P的值越小越好。
上图表示三种不同方法的雅克比行列式和采样网格情况,其中红色的部分表示行列式为负数的位置。在baseline模型中超过90%的不满意的位置都被消除了。所以有两种不同的方式可以减少形变场的重叠位置,同时保证配准的精度,一种是像Prob-VoxelMorph一样的直接方式,另一种是像循环一致性配准一样的间接方式。
上表是几种不同配准方法的 P P P和Dice的均值和标准差对比。