H-DenseUNet: Hybrid Densely Connected UNet for Liver and Tumor Segmentation from CT Volumes
这篇论文发表于2018年的IEEE Transactions on Medical Imaging,来自香港中文大学团队。
论文亮点
摘要
精准而智能的肝区域和肝肿瘤分割方法在临床实践中是高度需要的,它可以帮助医生进行肝癌诊断与治疗计划的确定。在最近的方法中,2D和3D FCN方法充当了立体图形分割的支柱。然而,2D卷积无法利用第三个维度的空间信息,同时,3D卷积则存在高昂的计算代价和GPU内存消耗的问题。为了解决这些问题,我们提出了一种新颖的混合密集连接的U-Net模型(H-DenseUNet),它由一个用于有效提取slice之间的特征的2D DenseUNet模型和一个基于肝脏肿瘤分割的自动上下文算法的分层融合三维上下文的3D DenseUNet。我们将H-DenseUNet的学习过程以一种端到端的方式公式化,其中切片之间的表征和切片内部的特征合一通过一个混合特征融合层联合优化。我们在MICCAI2017的肝分割挑战赛数据集和3DIRCADb数据集上进行了广泛地验证。我们的方法在肿瘤的分割结果上超越了其他SOAT方法,并且甚至在单个模型上,也取得了非常具有竞争力的表现。
介绍
从CT图像中对肝癌做出准确诊断是十分重要的。传统上,肝区域和肝病变区域都是通过医生一个切片一个切片地进行描述的,然而这不仅耗时,而且受到医生主观因素的影响。因此,自动的肝区域分割和肝病变(肿瘤)区域分割在临床应用中是十分必要的。
从CT图像中准确地进行分割出肝区域是非常困难的,而分割出肝肿瘤区域就更困难了。
为了解决这些困难,已经提出了许多分割方法,包括灰度阈值、区域增长和可变模型等。然而这些传统方法依赖于手工制作的特征,并且特征表达能力有限。最近,全卷积神经网络在识别问题上取得了很大的成功。
许多研究者在肝区域分割和肝肿瘤分割上使用了深度学习方法,这些方法大概可以分为两类
:
【1】
2D FCNs
,例如U-Net结构,多通道的FCN等;
【2】
3D FCNs
,用3D卷积代替2D卷积,并且以三维数据作为输入。
在临床诊断中,有经验的医生通常会根据
z轴上的相邻切片
来观察和分割肿瘤。然而,基于2D FCN的方法忽视了z轴上的上下文信息,这会导致分割精度的下降。具体来说,有些方法使用单张或者三张相邻的slice作为2D FCN的输入,随后将得到的2D 分割图堆叠成三维分割结果。尽管采取三张相邻slices作为输入的方法利用了相邻切片的信息,但是它们仍然不足以提取空间维度上的空间信息,这会使分割结果下降。为了解决这个问题,一些研究者也提出了使用多个平面切片的方法、以及RNN的方法来提取3D信息。相比于2D FCN,3D FCN的计算代价和GPU内存消耗过高。过高的GPU内存消耗也限制了网络的深度和滤波器的大小,而这二者是网络性能的关键。3D卷积沉重的计算代价也妨碍了它应用在大尺度的数据集上。
为了解决上述问题,我们提出了一种新颖的端到端系统,称之为混合密集连接型U-Net(H-DenseUNet),其中intra-slice(切片之间的)特征和3D 上下文信息都很有效地进行了提取,并且进行了联合优化以实现准确的肝区域和病变区域分割。我们的H-DenseUNet有以下两个技术上的成就:
深而高效的网络
:首先,为了完全地提取切片间的高层特征,我们设计了一个非常深而高效的网络,称之为2D DenseUNet,将密集连接路径和UNet连接的优势进行融合。密集连接路径来源于DenseNet,它解决了训练非常深的网络的困难。不同于DenseNet,我们增加了UNet连接,即在我们架构编码路径和解码路径之间的大范围跳跃连接。因此,网络可以保留低层空间特征来更好地探索切片之间的上下文信息。
混合特征探索
:第二,为了探索三维特征表达,我们设计了一种端到端的训练系统,称之为H-DenseUNet,其中切片间和切片内的特征都被有效地提取,并且联合地通过混合特征融合层(HFF)进行优化。具体而言,3D DenseUNet是由
2D DenseUNet通过自动上下文机制组合
而成的。有了来自2D DenseUNet的语义概率的指导,在3D DenseUNet上的优化扰动被很好地消除,这有利于3D 上下文信息提取。此外,在端到端系统下,由三维特征和有代表性的高层切片间特征组成的混合特征可以自动地融合,并一起进行优化以更好地进行肝和肿瘤识别。总的来说,这项工作有以下成就:
【1】我们设计了一个DenseUNet来有效地提取切片间的分层特征以进行肝和肿瘤分割,其中对密集连接路径和UNet连接进行了融合来提升肝肿瘤的分割性能。
【2】我们提出了一种H-DenseUNet框架来探索对于肝区域分割和肿瘤分割的混合特征(切片间特征和切片内特征)。这个混合特征学习框架很好地回避了2D网络会忽略三维上下文信息和3D网络存在很大计算代价的问题,并且可以作为一种新的探索3D上下文的新范例。
【3】我们的方法取得了极好的结果。
方法
上图展示了我们提出的肝区域和肝病变分割方法的流程。我们采用了级联学习策略来减少总的运算时间,它已经在许多识别任务中使用。首先,一个简单的ResNet框架被训练为获取一个快速但粗略的肝分割(
注意,这里未在图中画出!!
)。在ResNet提取ROI区域后,我们提出的H-DenseUNet可以通过一个2D DenseUNet 和3D DenseUNet,有效地探索切片内和切片间的特征,并通过混合特征融合层联合地对混合特征进行优化来实现精确肝区域和病变区域分割。
用于切片间特征提取的深度2D DenseUNet
切片间的特征提取路径遵循了DenseNet-161(即DenseNet论文提出的结构)的结构,它由重复的带有不同输出维度的密集连接块组成。在每一个密集连接块中,从任何一层到后面的子层中都有直接的连接,如上图
(c)
中所示。每一层都能产生k个特征图,k被称之为增长率。层之间的密集连接的一个好处在于它相比传统网络有更少的输出维度,避免学习大量的特征。此外,密集连接路径还能确保层之间的最大信息流动,这可以提升梯度的流动性,并且因此缓解在一个非常深的神经网络搜索全局最优时的扰动。
然而,原始的DenseNet-161用于目标分类任务。此外,由于分割任务的FCN网络常常需要采取池化或上采样操作,这可能会导致低层特征(即高分辨率)的信息丢失。综合这两种考虑,我们提出了一种2D DenseUNet,其中结合了密集连接路径和UNet跳跃连接路径的优点。具体来说,层之间的密集连接应用在每个小块中,来确保最大信息流动。同时,UNet在编码器和解码器之间的大范围跳跃连接用于保留低层信息。
用
I
∈
R
n
×
224
×
224
×
12
×
1
\mathbf{I} \in R^{n \times 224 \times 224 \times 12 \times 1}
I
∈
R
n
×
2
2
4
×
2
2
4
×
1
2
×
1
表示输入训练样本(对应于
224
×
224
×
12
224 \times224\times12
2
2
4
×
2
2
4
×
1
2
的输入三维数据),带有
Y
∈
R
n
×
224
×
224
×
12
×
1
\mathbf{Y} \in R^{n \times 224 \times 224 \times 12 \times 1}
Y
∈
R
n
×
2
2
4
×
2
2
4
×
1
2
×
1
的真实标签,其中
n
n
n
表示输入训练样本的batch大小,最后一个维度代表channel。分割标签中的
Y
i
,
j
,
k
=
c
\mathbf{Y}_{i, j, k}=c
Y
i
,
j
,
k
=
c
表示每一个像素
(
i
,
j
,
k
)
(i,j,k)
(
i
,
j
,
k
)
都被打上了
c
c
c
类的标签(背景、肝区域、肿瘤区域)。
用函数
F
\mathscr{F}
F
来表示从三维数据到三个相邻切片(即二维数据)的转换
。具体而言,沿着z轴每三个相邻的切片组合在一起,组数被转换到batch维上。即相当于如果
I
2
d
=
F
(
I
)
\mathbf{I}_{2 \mathbf{d}}=\mathscr{F}(\mathbf{I})
I
2
d
=
F
(
I
)
,则
I
2
d
∈
R
12
n
×
224
×
224
×
3
\mathbf{I}_{2 \mathrm{d}} \in R^{12 n \times 224 \times 224 \times 3}
I
2
d
∈
R
1
2
n
×
2
2
4
×
2
2
4
×
3
(
其中,
n
×
224
×
224
×
12
×
1
n\times224 \times224\times12\times1
n
×
2
2
4
×
2
2
4
×
1
2
×
1
的立体输入数据中(batch为n),如果将z轴上12个切片中每三个相邻的切片组合起来(有额外的padding),转换为二维数据则共有12n个组合,作为batch维,3相当于channel维
)表示2D DenseUNet的输入样本,如上图
(d)
中所示。
由于这个转换存在,2D 和 3D的 DenseUNet可以联合训练。
F
−
1
\mathscr{F}^{-1}
F
−
1
用于表示从三个相邻切片到三维数据的逆变换。
2D DenseUNet的详细描述和细节分别展示在上图
(c)
中和下表中。2D DenseUNet的的深度扩展到了167层,包括了卷积层、池化层、密集连接块、转移层和上采样层。密集块代表了几个小块的级联,其中所有层都进行了直接连接,如网络结构图
(c)
所示。为了改变特征图的大小,应用了转移层,它包含了一个batch normalization层和一个
1
×
1
1\times1
1
×
1
的卷积层+一个平均池化层。在转移层中包含了一个压缩因子,用于压缩特征图的数量,防止特征图的扩张(在我们的实验中被设为0.5)。上采样层通过双线性变换完成,在之后跟随着UNet的低层特征连接与一个3×3的卷积层。在每个卷积层之前,Batch Normalization方法和ReLU被应用于架构之中。
用于混合特征探索的H-DenseUNet
2D卷积会忽视空间信息,3D卷积则会带来高昂的计算代价。为了解决这些问题,我们提出了H-DenseUNet来共同地融合和优化切片间的特征和切片内的特征(切片间的特征可以理解为3D 上下文信息,即2D 卷积无法获取的特征,而切片内的特征可以通过2D卷积获取),以实现更好的肝区域分割和肿瘤分割。
为了融合从2D网络和3D网络提取的混合特征,特征的大小应该对齐。因此,从2D DenseUNet提取的2D特征图和分割图应该通过
F
−
1
\mathscr{F}^{-1}
F
−
1
被转换为3D形状。
然后,3D DenseUNet通过将原始的三维CT图像和2D网络提取出的分割图进行级联作为输入,从3D上下文中提取视觉特征。具体来说,在3D网络中的检测器不仅仅基于3D原始图像提取的特征训练,也基于从2D DenseUNet得到的分割概率图。在2D 分割结果的指导下,3D DenseUNet的训练与学习会更加高效。
将2D网络与 3D网络提取的特征相加后,得到混合特征
。这个混合特征将在HFF层进行联合学习和优化。
最终,3D 网络只需要9个小时就能收敛,而如果使用原始数据集去训练3D网络,则需要63个小时才能收敛,这证明了利用混合特征很好地减少了训练的计算代价。
损失函数、训练过程与测试过程
使用了加权交叉熵损失作为损失函数。
训练过程
:首先按照之前一篇论文的方式对ResNet(哪里来的ResNet)进行训练,得到一个粗略的肝分割结果。在2D DenseUNet中的编码器权重按照DenseNet提出论文训练结果的参数进行初始化,而2D DenseUNet的解码器则进行随机初始化。由于解码器的权重是随机初始化的,在训练的前几个epoch我们先不添加UNet的跳跃连接。在几个epoch后,UNet的跳跃连接被添加以共同训练模型。
为了更有效地训练H-DenseUNet,我们在使用的数据集上,首先通过交叉熵损失训练2D DenseUNet。随后,我们固定2D DenseUNet的参数,并且专注于通过交叉熵损失训练3D DenseUNet和HFF层,它们的权重都是随机初始化的。最后,整个网络通过下面的联合损失进行微调(fine-tune):
L
total
=
λ
L
(
y
,
y
2
^
d
′
)
+
L
(
y
,
y
H
^
)
L_{\text {total}}=\lambda L\left(y, y_{\hat{2} d}^{\prime}\right)+L(y, \hat{y_{H}})
L
total
=
λ
L
(
y
,
y
2
^
d
′
)
+
L
(
y
,
y
H
^
)
推理过程
:在测试阶段,我们首先得到一个粗略的肝分割结果(通过ResNet)。然后H-DenseUNet可以在肝分割结果ROI上生成精确的肝区域和肿瘤区域分割结果。
实验与结果
在验证集上的2D DenseUNet和H-DenseUNet的分割结果示例。红色区域表示分割的肝区域,绿色区域表示病变区域。灰色部分表示真实的肝区域,白色部分表示真实的病变区域。