“挤压和刺激”引导的Few-shot体图像分割(Medical Image Analysis 2020)
摘要
我们提出了一种新颖的Few-Shot框架,用于只有少量注释切片的体医学图像分割。与其它计算机视觉工作相比,主要的挑战在于缺乏预先训练的网络和医学扫描的体性质。为了应对这些挑战,我们提出了一种包含“挤压&刺激”块的Few-Shot分割新架构。我们的双臂架构由包含一个条件分支,它处理带注释的支持输入并生成特定于任务的特征表示。这个表示被传递给使用这个信息来分割新的查询图像的分割分支。为了促进条件分支和分割分支之间的有效交互,我们提出使用“通道挤压和空间激励”模块——一个轻量级的计算模块——使两个分支之间的交互作用更强,而模型复杂性的增加可以忽略不计。这一贡献使我们能够在不依赖于预先训练的模型的情况下进行图像分割,而预训练在医学扫描中通常是不可用的。此外,我们提出了一种有效的体分割策略,通过对支持体的几个切片与查询体的所有切片进行最佳匹配来提升性能。
存在的问题及解决方案
现在的Few-Shot分割大多针对于2D图像,且使用预训练来获取强力的先验知识。因此条件分支和分割分支之间的弱交互仍然有效。但医学图像缺乏有效的预训练模型,将Few-Shot的方案直接使用到医学图像上比较困难,条件分支和分割分支都得重新训练。然而,在缺乏预训练模型的情况下,弱交互将导致模型失稳甚至崩溃。
不同于一般的弱交互模式,我们提出了一种在条件分支和分割分支之间在多个位置上的强交互。强交互促进了有效的跨分支梯度流动,使得双分支训练变得更轻松,而不需要任何预训练模型。为此,我们提出一种通道挤压和空间激励(channel squeeze & spatial excitation)的模块sSE。这个块以条件分支的输出特征图作为输入,执行通道挤压操作来学习一个空间映射,这个空间映射被用于待分割图像特征图的空间激励。我们在所有的编码器、bottleneck和解码器块之间使用sSE块。SE块非常适合于实现分支之间的相互作用,因为它们权重参数少,因此只略微增加模型的复杂性。尽管其权重参数很少,但它们可以通过重新校准对待分割图像的特征产生强烈的影响。
现有的Few-Shot分割工作主要集中在2D图像上,而我们正在处理体医学扫描图像。在三维图像的所有切片上手动标注器官是一项耗时的工作。按照Few-Shot学习的思想,注释应该发生在几个稀疏选择的切片上。为此,我们提出了一种体分割策略,将支持体的几个带注释的切片与查询体的所有切片进行适当的配对,以保持分割的切片间的一致性。
方法
条件分支(Conditioner Arm)
条件分支的作用在于融合标注 L s L_s Ls支持图像 I s I_s Is的的视觉信息,生成特定任务(task-specific)的特征图,捕捉查询图像 I q I_q Iq中需要分割的内容。我们将条件分支的中间特征 图成为任务表示( task representation)。我们通过堆叠 I s I_s Is和二进制映射 L s ( α ) L_s(\alpha) Ls(α)来提供一个输入到条件分支的2通道输入,其中 L s ( α ) L_s(\alpha) Ls(α)表示支持图像掩码, α \alpha α表示语义类。这与经典的OSLSM中利用相乘来产生输入的思路相反,他们的动机是抑制背景像素,这样条件分支就可以专注于物体内部的图案(就像猫类的眼睛或鼻子图案)。但这并不适用于医学扫描图像,因为一个器官语义类中只有有限的纹理模式。例如,肝脏内的体素强度相当均匀,但边缘模糊。因此,我们将支持集的两部分都提供给网络,让它学习最优融合,从而提供查询图像的最佳分割。
整体网络架构如下图所示:
使用“挤压&激励”模块的交互块(Interaction Block using ‘Squeeze & Excitation ’ modules)
交互块在Few-Shot分割框架中起着关键作用,这些块使用条件分支得到的任务表示(task representation)作为输入,将它们传递给分割分支,以引导对查询图像的分割。理想情况下,这些block应该是轻量级的,只会略微增加模型的复杂性和计算时间,同时通过改善梯度流,可以方便网络的训练。
sSE说白了就是找到支持特征的各个特征图的重要程度(通道压缩),并将这一特征图重要程度转移到查询特征上来指导(在这里被说成是激励与校准)分割。
sSE块沿着通道压缩一个条件分支特征图
U
c
o
n
∈
R
H
×
W
×
C
′
U_{con}\in \mathbb{R}^{H\times W\times C\prime}
Ucon∈RH×W×C′,激活相应的分割分支特征图
U
s
e
g
∈
R
H
×
W
×
C
U_{seg}\in \mathbb{R}^{H\times W\times C}
Useg∈RH×W×C,传递来自支持集的信息,以辅助查询图像的分割。这里,我们考虑一种特殊的切片策略来表示输入张量
U
con
=
[
u
con
1
,
1
,
u
con
1
,
2
…
,
u
con
j
,
ι
,
…
,
u
con
H
,
W
]
\mathbf{U}_{\text {con }}=\left[\mathbf{u}_{\text {con }}^{1,1}, \mathbf{u}_{\text {con }}^{1,2} \ldots, \mathbf{u}_{\text {con }}^{j, \iota}, \ldots, \mathbf{u}_{\text {con }}^{H, W}\right]
Ucon =[ucon 1,1,ucon 1,2…,ucon j,ι,…,ucon H,W],其中
u
con
j
,
ι
∈
R
1
×
1
×
C
′
\mathbf{u}_{\text {con }}^{j, \iota} \in \mathbb{R}^{1 \times 1 \times C^{\prime}}
ucon j,ι∈R1×1×C′,对分割分支特征图做同样的切片
U
seg
=
[
u
seg
1
,
1
,
u
seg
1
,
2
…
,
u
seg
j
,
ι
,
…
,
u
seg
H
,
W
]
\mathbf{U}_{\text {seg }}=\left[\mathbf{u}_{\text {seg }}^{1,1}, \mathbf{u}_{\text {seg }}^{1,2} \ldots, \mathbf{u}_{\text {seg }}^{j, \iota}, \ldots, \mathbf{u}_{\text {seg }}^{H, W}\right]
Useg =[useg 1,1,useg 1,2…,useg j,ι,…,useg H,W]。空间压缩操作是使用卷积来执行的
q
=
W
s
q
⋆
U
c
o
n
\mathbf{q}=\mathbf{W}_{s q} \star \mathbf{U}_{\mathrm{con}}
q=Wsq⋆Ucon来执行的,其中
W
s
q
∈
R
1
×
1
×
C
′
\mathbf{W}_{s q} \in \mathbb{R}^{1 \times 1 \times C^{\prime}}
Wsq∈R1×1×C′,生成投影张量
q
∈
R
H
×
W
\mathbf{q} \in \mathbb{R}^{H \times W}
q∈RH×W,这个投影
q
\mathbf{q}
q通过一个sigmoid操作σ(·)来激活到[0,1],这是用来在空间上重新校准或激活
U
s
e
g
\mathbf{U}_{seg}
Useg来生成:
体分割策略(Volumetric Segmentation Strategy)
我们使用2D图像作为支持和查询来训练网络,但是在测试阶段,我们需要分割的是3D查询体。因此我们需要从支持体中选择一组稀疏的带注释的切片,由它们构成支持集。直接扩展对查询体的分割很难,因为没有现成的策略来将上述选择的支持切片和查询体的所有切片对应(这样才能在体中产生好的分割)。我们在这里提出一种解决策略。
假定提供了支持体中的 k k k个带注释的切片,查询体按照以下步骤进行分段:
- 给定一个语义类,我们首先指出该器官在支持体和查询体上所在的切片范围(沿着一个固定的方向)。假定查询体的范围是 [ S s , S e ] [S^s,S^e] [Ss,Se],支持体的范围是 [ Q s , Q e ] [Q^s,Q^e] [Qs,Qe]。这里的上标表示开始 s s s和结束 e e e切片索引。
- [ S s , S e ] [S^s,S^e] [Ss,Se]和 [ Q s , Q e ] [Q^s,Q^e] [Qs,Qe]被分为 k k k个平均间隔的切片组。分别用 [ { S 1 i } , … , { S k i } ] \left[\left\{S_{1}^{i}\right\}, \ldots,\left\{S_{k}^{i}\right\}\right] [{S1i},…,{Ski}]和 [ { Q 1 i } , … , { Q k i } ] \left[\left\{Q_{1}^{i}\right\}, \ldots,\left\{Q_{k}^{i}\right\}\right] [{Q1i},…,{Qki}]表示,这里的下标表示组号。
- 在每个 k k k支持体组中,中心切片 [ S 1 c , … , S k c ] \left[S_{1}^{c}, \ldots, S_{k}^{c}\right] [S1c,…,Skc]被标注作为支持集。
- 我们将带注释的中心切片 S j c S^c_j Sjc与组 { Q j i } \left\{Q_{j}^{i}\right\} {Qji}的所有切片配对,形成了条件分支和分割分支的输入。
实验结果