1. 简介
DeepIGeoS 是2019年发表在 TPAMI 中的一篇文章,与2016年发表在 CVPR 上的 Deep Interactive Object Selection 文章不同,本文提出的方法是在自动分割的结果上让用户对其进行标注,细化分割。
创新点:与现有的一些交互式分割方法相比,所需要的用户标注的工作量较少;响应的时间更快,每次用户交互后不再需要对全部模型进行重新训练,而是采取预训练+模拟用户交互的过程。
图 1,两阶段的交互式分割框架
DeepIGeoS 是一个两阶段的框架,使用了两个 CNN,如图1所示。Stage 1 中的 P-Net 用于初始分割(粗分割),将初始分割结果与用户的标注相结合,用作 Stage 2 中 R-Net 的输入,用于 refine 初始分割。其中 CRF-Net(f) 是我们提出的使用自由形式二元势的具有反向传播的 CRF,结合了用户交互作为强约束组成 CRF-Net(fu)。具有如下特性:
- 使用更强大的基础学习模型,即具有自动特征学习的 CNN 可以应付大量的训练集;
- 需要较少的用户交互,用户标注从高质量的自动分割图的基础上开始;
- 可以实时响应用户交互,比较高效,可以节省用户时间。
2. 方法
2.1 基于用户交互的测地线距离图
在 stage 1 初始分割图像的基础上,用户进行标记,标记分为前景和背景,相同标记的交互被转换为距离图。
测地距离图比欧式距离可以更好地反应图像中的边缘和上下文信息,提高同质区域的标签一致性
设 S f , S b S_f, S_b Sf,Sb 分别为前景和背景标记的像素集, i i i 为图像 I I I 的一个像素点,则 i i i 到标注集 S ( S ∈ { S f , S b } ) S\left(S\in\left\{S_f,\;S_b\right\}\right) S(S∈{Sf,Sb}) 的无符号测地线距离为:
其中
P
i
,
j
P_{i,\;j}
Pi,j 为像素
i
i
i 到
j
j
j 的所有路径的集合,
p
p
p 是一条可行的路径,由
S
∈
[
0
,
1
]
S\in\left[0,\;1\right]
S∈[0,1] 进行参数化。
u
(
s
)
=
p
′
/
∥
p
′
(
s
)
∥
\boldsymbol u\left(s\right)=p'/\left\|p'\left(s\right)\right\|
u(s)=p′/∥p′(s)∥ 是与路径方向相切的单位向量。
如果没有为前景或背景绘制标记,则相应的测地距离图将填充随机数。
图 2,用户标注转换为测地距离图的示例
如图2所示,(a) 中绿色曲线为初始分割结果,红色为前景标记,青涩为背景标记,(d)、(e) 分别为转换后的基于前景和背景交互的测地线距离图像。原始图像 (a)、初始分割 (c)、测地距离 (d)(e) 图像组合在一起,得到一个四通道的图像,用以作为第二个 CNN (R-Net) 的输入。
2.2 使用扩张卷积的保持分辨率的 CNN
R-Net 改编自 VGG16 (13个卷积层、3个全连接层、5个池化层)
图 3,含有 CRF-Net(f) 的用于 2D/3D 分割 CNN 结构
以图 (a) 为例,前13个卷积层被分为5个块。第1、2块分别有两个卷积层,其余每个块都有3个卷积层。卷积层的参数在深蓝色的矩形框中(包括内核尺寸、输出通道、膨胀参数)。在所有的卷积层中,卷积核的大小都固定为3×3。块1到块6是保留分辨率的。
每个卷积层的步长设置为1。每个块中卷积的输出通道数设置为常数C。为了使用多尺度特征,作者将来自不同块的特征连接起来得到一个长度为5C的组合特征。这个特征被输入到一个分类器中,该分类器由两个附加层实现,如图(a)中的块6所示。这两个层使用大小为1×1、扩张参数为0的卷积核。块6为每个像素提供属于前景或背景类的初始分值。
图 (b) 为用于3D分割的网络结构,其与2D结构类似,只不过为了减少3D图像的内存损耗,作者在图像输入时,使用了下采样,并将块1到块5的输出特征通过1×1×1的卷积块压缩了4倍。
为了获得感受野的指数级增长,VGG16在每个块之后都使用最大池化和下采样层,这种实现方式会以指数级降低特征图的分辨率。因此,为了保持分辨率,本文移除了最大池化层和下采样层,并在每个块中使用扩张卷积。
为了在给定标注时获得更空间一致性的分割并添加硬约束,作者在块6输出的基础上应用CRF。CRF由循环神经网络实现,可以与 P-Net 或 R-Net 联合训练。CRF-Net为每个像素提供正则化预测,该预测被输入交叉熵损失函数层。
2D/3D R-Net 使用与 2D/3D P-Net 相同的结构,只是它的输入通道数为 C 1 + 3 C_1+3 C1+3 (如上一小节所示),并且由于集成了用户交互,所以 CRF-Net(f) 变为了 CRF-Net(fu)。
2.3 具有自由形式的二元势函数和用户约束的反向传播 CRF-Net
CRF-Net 可以通过反向传播进行训练,本文并没有使用高斯函数,而是扩展了 CRF,使得二元势函数是自由形式的函数,称为 CRF-Net(f)。另外在交互式细化上下文中将用户的标注集成到 CRF-Net(f) 中,称为 CRF-Net(fu)。
即:CRF-Net(f) 连接到 P-Net,CRF-Net(fu) 连接到 R-Net。
设 X X X 是图像 I I I 中标签集 L = { 0 , 1 , . . . , L − 1 } L=\left\{0,\;1,\;...,\;L-1\right\} L={0,1,...,L−1} 的标签映射,吉布斯分布(Gibbs distribution)建模 CRF 中给定 I I I 的 X X X 的概率:
即为描述单像素所对应的标签的概率,其中
Z
(
I
)
Z\left(I\right)
Z(I) 是配分函数的归一化因子。
Eq. (6) 被称为吉布斯能量函数(Gibbs energy),由两部分组成。第一部分称为一元势函数,描述单点像素和标签的关系,本文中的一元势是从 P-Net 或 R-Net 中获得,给每个像素一个分类分(classification scores)。第二部分,二元势函数描述像素对之间的关系,相似的像素分配相同的标签,相差较大的分配不同的标签。最小化能量函数可以得到更为精确的分割。CRF 能使图片尽可能在边界处分割。
Eq. (7) 是二元势函数,其中
μ
(
x
i
,
x
j
)
\mu\left(x_i,\;x_j\right)
μ(xi,xj) 表示
i
,
j
i, j
i,j 标签之间的兼容性,由大小为
L
×
L
L×L
L×L 的矩阵表示。
f
(
⋅
)
f\left(\cdot\right)
f(⋅) 是关于特征差异
f
~
i
,
j
=
f
i
−
f
j
{\widetilde{\boldsymbol f}}_{i,\;j}={\boldsymbol f}_i\boldsymbol-{\boldsymbol f}_{\mathbf j}
f
i,j=fi−fj (其中
f
i
,
f
j
{\boldsymbol f_i}, {\boldsymbol f_j}
fi,fj 分别表示为
i
,
j
i, j
i,j 的特征向量)和欧式距离
d
i
,
j
d_{i, j}
di,j 的函数,本文并没有将其定义为高斯函数或多个高斯函数的组合,而是将其设置为由全连接神经网络(Pairwise-Net)表示的自由函数,如下图4所示,它可以在训练中学习,有两个隐藏层和一个输出层。
图 4,Pairwise-Net 结构
由于二元势规模越大,导致计算量就越大,所以将复杂的后验分布简化为 Q Q Q ,在最小散度(KL-divergence) D ( Q ∥ p ) D\left(Q\parallel p\right) D(Q∥p) 下对 Q Q Q 迭代更新,得到原始后验概率分布的近似解。
其中
L
L
L 是标签集,
i
,
j
i, j
i,j 是像素对。
对于提出的 CRF-Net(fu),在用户提供的标注集 S f b = S f ∪ S b S_{fb}=S_f\cup S_b Sfb=Sf∪Sb 的情况下,强制标注集中像素的概率为1或0,每次迭代的更新规则为:
其中
s
i
s_i
si 表示用户提供的像素
i
i
i 的标签,它位于
S
f
b
S_{fb}
Sfb 中。
本文使用 CRF-Net 中的实现,通过 RNN 中的多阶段平均场(平均近似场通常用于提升 CRF 的 inference 效率,同时允许通过反向传播学习参数 )法更新 Q Q Q。每个平均场分裂方程(Eq. 8)分为4个步骤,包括消息传递、兼容性转换、添加一元势和归一化。
2.4 实施细节
2.4.1 测地线距离的变换
光栅扫描算法用于计算测地线距离变换,方法是应用前向扫描和后向扫描,2D使用3×3内核,3D使用3×3×3内核。由于访问连续块中的图像内存,所以速度很快。对于本文提出的具有自由形势的二元势的 CRF-Net,没有使用整个图像中的全连接,而是使用基于局部块的像素连接。原因如下:
- 仅当二元势函数为高斯函数时,效率才会高
- 长距离依存(long-distance dependency)处理医学影像难度大。由于医学影像的对比度通常都很低,在这种情况下,长距离依存可能会导致目标像素的标签被大量具有相似外观的像素所影响
因此,为了保持良好的效率并避免长距离损失,本文定义了以该像素为中心的局部块内的一个像素的成对连接。2D图像的块大小设置为7×7,3D图像的块大小为5×5×3。
2.4.2 自由形式的二元势函数
Eq. (7) 中的 μ ( ⋅ ) \mu\left(\cdot\right) μ(⋅) 初始化为 μ ( x i , x j ) = [ x i ≠ x j ] \mu\left(x_i,\;x_j\right)=\left[x_i\neq x_j\right] μ(xi,xj)=[xi=xj],其中 [ ⋅ ] [\cdot] [⋅] 为艾佛森括号(Iverson bracket),满足条件则为1,否则为0。
图4中,第1和第2个隐藏层分别有32和16个神经元。该 Pairwise-Net 是由全卷积神经网络实现的,2D为1×1内核,3D为1×1×1内核。在预训练中使用对比度敏感函数(contrast sensitive function)的近似值初始化 Pairwise-Net:
其中
F
F
F 是特征向量
f
i
f_i
fi 和
f
j
f_j
fj 的维度,
ω
,
σ
\omega,\sigma
ω,σ 是分别控制函数大小和形状的参数。根据经验分别设置为0.5和0.8。
将 f i , f j f_i, f_j fi,fj 设置为 P-Net 的输入通道(即本文中的图像强度),用以简化实现并获得对比度敏感的二元势能(contrast-sensitive pairwise potentials)。
2.4.3 预训练 Pairwise-Net
生成一个包含100k样本的训练集 T ′ = { X ′ , Y ′ } T'=\left\{X',\;Y'\right\} T′={X′,Y′}
其中 X ′ X' X′ 是模拟连接 f ~ i j {\widetilde{\boldsymbol f}}_{ij} f ij 和 d i j d_ij dij 的特征集, Y ′ Y' Y′ 是模拟 Eq. (11) 的一组预测值集。对于 T ′ T' T′ 中的每个样本 s s s ,特征向量 x s ′ x'_s xs′ 的维度为 F + 1 F+1 F+1 ,其中前 F F F 个维度表示 f ~ i j {\widetilde{\boldsymbol f}}_{ij} f ij 的值,最后一个维度表示 d i j d_{ij} dij 。 x s ′ x'_s xs′ 的第 c c c 个通道用随机数 k ′ k' k′ 填充, k ′ ∼ N ( 0 , 2 ) k'\sim N\left(0,\;2\right) k′∼N(0,2) 表示 c ⩽ F c\leqslant F c⩽F , k ′ ∼ U ( 0 , 8 ) k'\sim U\left(0, \;8\right) k′∼U(0,8) 表示 c = F + 1 c=F+1 c=F+1 。
x s ′ x'_s xs′ 的预测值 y s ′ y'_s ys′ 的 ground truth 从 Eq. (11) 得到。在生成 X ′ X' X′ 和 Y ′ Y' Y′ 后用具有二次损失函数的随机梯度下降(SGD)算法来预训练 Pairwise-Net。
2.4.4 预处理
所有图像都通过训练集的平均值和标准差进行归一化。
通过垂直或水平翻转、角度范围为 [−𝜋/8,𝜋/8] 的随机旋转和 [0.8,1.25] 范围的随机缩放来实现数据增强。使用交叉熵损失函数和 SGD 算法进行优化, mini-batch 为 1, momentum 为 0.99,权重衰减为 5 × 1 0 − 4 5\times10^{-4} 5×10−4。每5k次迭代,学习率减半。
训练 P-Net
由于恰当的 P-Net 和 CRF-Net(f) 的初始化有助于联合训练的更快收敛,因此用 CRF-Net(f) 分三步训练 P-Net:
- 对 P-Net 进行预训练(使用初始学习率 1 0 − 3 10^{-3} 10−3 和最大迭代次数 100k)
- CRF-Net(f) 中的 Pairwise-Net 进行了如前所述的预训练
- P-Net 和 CRF-Net(f) 联合训练(初始学习率为 1 0 − 6 10^{-6} 10−6,最大迭代次数为50k)
训练 R-Net
训练完 P-Net 后,接下来训练 R-Net。该部分使用自动模拟用户交互来训练 R-Net。
首先,使用带有 CRF-Net(f) 的 P-Net 来获得每个训练图像的自动分割,它与 ground truth 进行比较以找到错误分割的区域。然后通过随机采样该区域中的 n 个像素来模拟每个错误分割区域上的用户交互。假设一个连接的欠分割或过度分割区域的大小为 𝑁 𝑚 𝑁_𝑚 Nm,如果 𝑁 𝑚 < 30 𝑁_𝑚<30 Nm<30 则将该区域的 n 设置为0,否则根据经验设置 [ 𝑁 𝑚 / 100 ] [𝑁_𝑚/100] [Nm/100]。
图 5,训练图像上的模拟用户交互示例。绿色为初始分割,黄色为 ground truth,红色为模拟欠分割区域的点击,青色为模拟过度分割区域的点击。
总结
DeepIGeoS 首先对 P-Net 进行预训练,然后用 P-Net 对训练数据的图像进行分割,分割后得到 y’ 与 y 进行比较来发现错误分割的区域,然后对错误区域进行标记,用于训练R-Net。
3. 实验
3.1 比较方法和评价指标
stage 1,将 P-Net 与 FCN 和 DeepLab 进行 2D 分割的比较,DeepMedic 和 HighRes3DNet 用于 3D 分割的比较。另外,将 CRF-Net(f) 与 Dense CRF, CRF-Net(g) 进行比较。
stage 2,比较了3种处理用户交互的方法:
- Min-cut user-editing,其中初始概率图(在本文中指 P-Net 的输出,初始分割图像)与用户交互相结合,以使用 min-cut 解决 𝐸(𝑥) 最小化问题
- 在 R-Net 中使用用户交互的欧氏距离,称为 R-Net(Euc)
- 所提出的具有用户交互互测地距离的 R-Net
另外,本文还将 DeepIGeoS 与其它几种交互式分割方法进行了比较。
2D:
- Geodesic Framework,它根据与用户提供的用于像素分类的标注的测地线距离计算概率
- Graph Cuts 将分割建模作为基于用户交互的最小切割问题
- Random Walks,根据 random walker 到达前景或背景种子的概率为像素分配标签
- SlicSeg 使用 Online Random Forests 从标注中学习并预测剩余的像素
3D:将 DeepIGeoS 与 GeoS、ITK-SNAP 进行了比较。
定量评估:Dice 系数和 ASSD(平均对称表面距离),都是常用的评价标准。
其中
𝑅
𝑎
𝑅_𝑎
Ra,
𝑅
𝑏
𝑅_𝑏
Rb 分别表示预测值的集合、真实值集合。
Dice系数是医学图像比赛中使用频率最高的度量指标,它是一种集合相似度度量指标,通常用于计算两个样本的相似度。在医学图像中经常用于图像分割,阈值范围为[0, 1],1最好,0最差。
其中,
𝑆
𝑎
𝑆_𝑎
Sa,
𝑆
𝑏
𝑆_𝑏
Sb 分别表示预测值和真实值的目标表面点的集合,
𝑑
(
𝑖
,
𝑆
𝑏
)
𝑑(𝑖,𝑆_𝑏 )
d(i,Sb) 是
𝑖
𝑖
i 和
𝑆
𝑏
𝑆_𝑏
Sb 之间的最短欧氏距离。
ASSD即预测表面点集中所有点到参考表面点集的最小欧氏距离,并计算所有距离的平均值,0为最佳,INF(infinity)为最差。
3.2 胎儿 MRI 的 2D 胎盘分割
3.2.1 stage 1: 2D P-Net 与 CRF-Net(f) 的自动分割
图 6,2D胎盘分割。不同分割网络的分割结果
第3行的 (b5) 表示只使用了块5的特征,而没有使用组合特征(图3所示)。最后一行显示了DeepIGeoS的交互式优化结果。其中ground truth由经验丰富的放射科医生手动描绘。
表 1,不同网络和 CRF 的定量比较
CRF-Net(g) 的二元势约束为高斯函数,CRF-Net(f)是本文提出的二元势约束为自由形式的CRF,其可在训练中学习。
图 7,基于 P-Net,添加了不同 CRF 的分割结果
3.2.2 stage 2: 2D R-Net 与 CRF-Net(fu) 的交互式 refinement
图 8,不同的分割效果。第一列为基于 2D P-Net 和 CRF-Net(fu) 的初始分割结果,将其作为输入。白色箭头显示局部细节的差异,红色是用户标记的前景,青色是背景标记
3.2.3 与其它交互方法的比较
图 9,与其它交互式分割的比较
须说明的是,除了本文的方法外,其它的方法在第一次分割前必须先进行标注。在 refinment 阶段,DeepIGeoS 只需要两次标注,其它方法需要更多的标注才能得到相似的结果。从图9可以看出,本文的方法所需要的用户工作量要少的多。
图 10,定量结果的比较
3.3 来自 FLAIR Images 的 3D 脑肿瘤分割
Flair也是一种成像技术,能很好的表现肿瘤部位周遭情况,清晰的表现出浮肿区域。
本部分使用了2015年脑肿瘤分割挑战赛(BraTS)训练集,其中包含274个病例的图像。几位专家手动勾画出了正确的范围。
本文随机选择234个案例进行训练,其余40个案例进行测试。所有这些图像都被剥去头骨并重新采样到尺寸为240×240×155。3D P-Net和R-Net的特征通道数为C=16。
3.3.1 使用 CRF-Net(f) 的 3D P-Net 自动分割
图 11,不同分割网络的分割结果
表 3,定量比较
图 12,P-Net 结合不同 CRF 的分割结果
3.3.2 R-Net 与 CRF-Net(fu) 的交互式细化
图 13,R-Net 分割结果的比较
表 4,定量比较
3.3.3 与其它 3D 交互方法的比较
图 14,与其它交互式分割方法的比较
第一行由于对比度高,所以效果都不错。第二行的对比度地,所以区别比较明显。
图 15,定量比较
4. 结论
- 本文提出了一个用于 2D 和 3D 医学图像分割的基于深度学习的交互式框架。
- 该框架包含两个阶段:使用 P-Net 获得初始自动分割结果,然后将用户的标注交互转换为测地距离图,集成到 R-Net 中用以改进初始分割结果。
- 本文还提出了一种用于密集预测(dense prediction)的具有扩张卷积的分辨率保持网络结构,并扩展了现有的基于 RNN 的 CRF,使其可以学习自由形式的二元势并利用用户交互作为硬约束。
- 通过实验部分,证明了本文提出的方法比自动分割的 CNN 取得了更好的结果。与传统的交互方法相比,用户花费的时间要少得多,并且可以实现更高的 3D 脑肿瘤分割精度。
- 该框架将来可以扩展到处理多个器官。