利用解剖相关性模拟One-Shot推理的泛化器官分割模型( IPMI 2021)
摘要
模仿学习是人类最重要的能力之一,在人类的计算神经系统中起着至关重要的作用。在医学图像分析中,给定几个样本(锚),有经验的放射科医生有能力通过模仿从现有器官类型中学习到的推理过程来描绘不熟悉的器官。受此启发,我们提出了OrganNet,它从一组标注的器官类别中学习一个广义的器官概念,然后将这个概念转移到没见过的类别中。在本文中,我们证明了这种过程可以集成到One-Shot分割任务中。我们提出金字塔推理模块(PRMs)来建模锚点和目标体(volume)之间的解剖相关性。在实践中,该模块首先计算目标和锚点CT体之间的相关矩阵。然后,利用该矩阵对锚定体及其分割掩码的特征表示进行变换。最后,OrganNet学习融合来自不同输入的表示,并预测目标体的分割结果。大量实验表明,OrganNet能够有效地抵御器官形态的巨大变化,并在One-Shot分割任务中取得了最新的结果。此外,即使与完全监督的分割模型相比,OrganNet仍然能够产生令人满意的分割结果。
存在的问题及解决方案
大量文献在深度学习框架内利用解剖相关性进行医学图像分割。解剖相关性也是基于图谱(atlas-based)分割的基础,将一幅或多幅标注过的参考图像(即图谱)基于解剖相似性非刚性配准到目标图像,并将图谱的标签传播到目标图像作为分割输出。与这些方法不同的是,OrganNet有能力通过使用推理过程来学习图像之间的解剖相似性。通过这种设计,提出的方法能够学习一个广义器官概念来进行One-Shot分割。
Zhao等人提出了一种自动数据增强方法,用于在磁共振成像(MRI)脑扫描上合成标记医学图像。但是,当锚点和目标图像之间只有很小的变化时,DataAug被严格限制在对象分割上。Roy等人通过将支持体的几个切片与查询体的所有切片进行最优配对,提出了一种Few-Shot体分割器。本文侧重于一个更现实的设置:我们使用丰富注释的器官来辅助分割较少注释的器官。此外,我们的方法与孪生学习(siamese learning)部分相关,通常采用一对图像作为输入。
在训练阶段,给定一幅目标图像,对锚点图像及其标注掩码进行配对。首先进行与锚点图像与目标图像的配准。然后,提出的推理模块试图建模锚点和目标图像之间的解剖相关性,之后利用相关矩阵对锚点图像及其对应的掩码特征表示进行变换,并与目标图像特征共同学习得到最终的预测结果。
方法
如图2所示,OrganNet有三个编码器,分别学习target volume、anchor volume和anchor mask表示。此外,我们提一种金字塔推理模块(PRMs)将解码器分为两部分:一部分学习anchor volume及其配对的anchor mask(前两个分支)的表示,另一个用于构建target volume和anchor volume之间的解剖相似性(底部两个分支)。
该OrganNet是建立在经典的3D U-Net结构上,我们将其扩展为一个三编码器版本,以包括额外的监督锚图像及其注释。在实践中,我们将OrganNet设计为轻量级的,以缓解医学图像分析中由于小数据集而引起的过拟合问题。
Pyramid Reasoning Modules(PRMs)
这些模块的设计是为了解决目标体和锚体之间器官形态和结构呈现不同程度变化的情况。由于不同特征金字塔中的特征捕获了多尺度信息,我们提出用一个推理函数对每个金字塔层的信息进行聚合。为了考虑两幅图像之间的位移,在浅层采用大尺寸的邻域,而在深层采用小尺寸的邻域。这种分配的根本原因是浅层的接受域比深层的接受域小。具体而言,首先计算目标体与锚体特征映射的相关矩阵。然后,我们利用这个矩阵分别变换锚输入的特征表示和它的分割掩码。最后,我们将三个输入的表示连接起来,并将它们作为下一层的输入。
如图3所示,在第
l
l
l层,每个推理模块有三个输入张量
X
l
,
t
X^{l,t}
Xl,t,
X
l
,
a
X^{l,a}
Xl,a,
X
l
,
m
X^{l,m}
Xl,m,对应于target volume, anchor volume和anchor mask。我们首先对以上三个输入张量应用三个1 × 1 × 1卷积算子,使它们归一化到相同的尺度,输出
X
h
,
t
X^{h,t}
Xh,t,
X
h
,
a
X^{h,a}
Xh,a,
X
h
,
m
X^{h,m}
Xh,m,大小为
H
l
×
W
l
×
D
l
×
M
l
H^{l} \times W^{l} \times D^{l} \times M^{l}
Hl×Wl×Dl×Ml。
为了对
X
h
,
t
X^{h,t}
Xh,t和
X
h
,
a
X^{h,a}
Xh,a之间的相关性进行建模,我们对两个张量的邻近区域应用内积运算。具体来说,给定一个
X
h
,
t
X^{h,t}
Xh,t中的向量
X
i
,
j
,
k
h
,
t
X_{i, j, k}^{h, t}
Xi,j,kh,t,其中
0
≤
i
<
H
l
,
0
≤
j
<
W
l
0 \leq i<H^{l}, 0 \leq j<W^{l}
0≤i<Hl,0≤j<Wl且
0
≤
k
<
D
l
0 \leq k<D^{l}
0≤k<Dl。我们根据
X
h
,
a
X^{h,a}
Xh,a的邻接区域构建内积:
其中
{
i
′
,
j
′
,
k
′
}
∈
{
−
⌊
s
2
⌋
,
…
,
⌊
s
2
⌋
}
\left\{i^{\prime}, j^{\prime}, k^{\prime}\right\} \in\left\{-\left\lfloor\frac{s}{2}\right\rfloor, \ldots,\left\lfloor\frac{s}{2}\right\rfloor\right\}
{i′,j′,k′}∈{−⌊2s⌋,…,⌊2s⌋}。
s
s
s表示相邻区域的大小,随着层深的变化而变化。
n
n
n表示总索引,可以计算为:
得到的
X
o
∈
R
s
3
×
H
l
×
W
l
×
D
l
X^{o} \in \mathcal{R}^{s^{3} \times H^{l} \times W^{l} \times D^{l}}
Xo∈Rs3×Hl×Wl×Dl。这一张量即用于表达
X
h
,
t
X^{h,t}
Xh,t和
X
h
,
a
X^{h,a}
Xh,a之间的解剖相似性。然后,我们对
X
o
X^{o}
Xo的第一维数应用softmax归一化,并将其扩展到
s
3
×
H
l
W
l
D
l
M
l
s^{3} \times H^{l} W^{l} D^{l} M^{l}
s3×HlWlDlMl:
考虑到矩阵乘法的效率,引入了im2col运算来将
X
h
,
a
X^{h,a}
Xh,a和
X
h
,
m
X^{h,m}
Xh,m转换成大小
s
3
×
H
l
W
l
D
l
M
l
s^{3} \times H^{l} W^{l} D^{l} M^{l}
s3×HlWlDlMl的张量。这样,我们可以用Hadamard积将它们与
X
w
X^w
Xw相乘,并应用求和来聚合相邻滤波器的贡献。形式上,计算过程可以形式化为:
其中
{
X
p
,
X
q
}
∈
R
H
l
×
W
l
×
D
l
×
M
l
\left\{X^{p}, X^{q}\right\} \in \mathcal{R}^{H^{l} \times W^{l} \times D^{l} \times M^{l}}
{Xp,Xq}∈RHl×Wl×Dl×Ml,
s
u
m
sum
sum应用于第一维。最后,我们连接三个输出来形成下一层的输入:
一般来说,推理模块学习如何从看到的器官中对齐锚点和目标体的语义表示。在推理阶段,学习到的规则可以很好地应用于使用one-shot对作为监督信号的未知类。
实验结果
总结
在这篇论文中,我们提出了一种新的One-Shot医学分割方法,通过锚点和目标CT体之间的推理,使我们能够学习广义器官概念。所提出的器官网在一个单一的深度学习框架内对关键组成部分(即图像之间的解剖相似性)进行建模,大量的实验证明了所提出的OrganNet的有效性。