摘要
本文提出了一种新的、功能更强大的医学图像分割架构UNet++。我们的架构本质上是一个深监督的编码器-解码器网络,其中编码器和解码器子网络通过一系列嵌套的、密集的跳跃路径连接。重新设计的跳跃连接旨在减少编码器子网络和解码器子网络的特征图之间的语义鸿沟。我们认为,当解码器和编码器网络中的特征映射在语义上相似时,优化器的学习任务将更加容易。我们评估了UNet++与U-Net和wide U-Net架构在多种医学图像分割任务中的应用:低剂量胸部CT扫描中的结节分割,显微镜图像中的核分割和结肠镜检查视频中的息肉分割。实验表明,相比U-Net和wide U-Net上,深度监管的UNet++平均IoU增益分别为3.9和3.4点。
介绍
最新的图像分割模型是编码器-解码器体系结构的变体,例如U-Net [9]和全卷积网络(FCN)[8]。 这些用于分割的编码器/解码器网络具有一个关键的相似性:跳跃连接,它将来自解码器子网络的深层,语义,粗粒度的特征图与来自编码器子网络的浅层,低层,细粒度特征图结合在一起。跳跃连接已经被证明可以有效地恢复目标对象的细粒度细节;即使在复杂的背景下,也可以生成具有精细细节的分割mask。 跳跃连接也是成功实现实例分割模型(例如Mask-RCNN)的基础,该模型可以对被遮挡的对象进行分割。 可以说,自然图像中的图像分割已达到令人满意的性能水平,但是这些模型是否满足医学图像的严格分割要求?
在医学图像中对病变或异常进行分割比在自然图像中需要更高的准确性。 尽管精确的分割mask在自然图像中可能并不重要,但即使医学图像中的边缘分割错误也可能导致临床环境中的用户体验不佳。 例如,结节周围细微的针状体可能表明结节是恶性的。 因此,从分割mask中排除它们会从临床角度降低模型的可信度。 此外,不正确的分割也可能导致后续计算机生成的诊断发生重大变化。 例如,在纵向研究中对结节生长的错误测量可能会导致将不正确的Lung-RADS类别分配给筛查患者。 因此,期望设计出可以有效地恢复医学图像中目标对象的精细细节的更有效的图像分割架构。
为了满足对医学图像中更精确的分割的需求,我们提出了UNet ++,这是一种基于嵌套和密集跳跃连接的新分割架构。我们架构背后的基本假设是,在与来自解码器的相应语义丰富的特征图融合之前,逐渐丰富来自编码器网络的高分辨率特征图,该模型可以更有效地捕获前景对象的细粒度细节。我们认为,当来自解码器和编码器网络的特征映射在语义上相似时,网络的学习任务将会变得更加轻松。这与U-Net中通常使用的普通跳跃连接形成对比,后者将高分辨率特征图从编码器直接转发到解码器网络,从而导致语义上不相似的特征图的融合。根据我们的实验,我们建议的体系结构是有效的,与U-Net和广泛的U-Net相比,可显著提高性能。
相关工作
Long等人[8]首先引入了全卷积网络(FCN),而U-Net由Ronneberger等[9]人引入。它们都有一个关键的思想:跳跃连接。在FCN中,将上采样的特征图与从编码器跳跃的特征图相加,而U-Net将它们连接(concat)起来,并在每个上采样步骤之间添加卷积和非线性。跳跃连接已显示有助于在网络输出处恢复完整的空间分辨率,从而使全卷积方法适用于语义分割。 受DenseNet架构的启发[5],Li等人[7]提出H-denseunet用于肝和肝肿瘤的分割。本着相同的精神,Drozdzal等人[2]系统地研究了跳跃连接的重要性,并在编码器中引入了短跳跃连接。 尽管上述架构之间的细微差别,但它们都倾向于融合编码器和解码器子网络在语义上不同的特征图,根据我们的实验,这可能会降低分割性能。
最近另外两个相关的工作是GridNet [3]和Mask-RCNN [4]。GridNet是一种编码器-解码器体系结构,其中的特征图以网格形式进行布线,推广了几种经典的分割架构。 但是,GridNet在跳跃连接之间缺少上采样层。 因此,它不代表UNet ++。 Mask-RCNN可能是用于目标检测,分类和分割的最重要的元框架。我们想指出的是,只需用建议的嵌套密集跳跃连接替换普通的跳跃连接,就可以将UNet ++轻松部署为Mask-RCNN中的骨干架构。 由于篇幅所限,我们无法将Mask RCNN的结果以UNet ++作为主干架构。 但是,有兴趣的读者可以参考补充材料以获取更多详细信息。
图1:(a)UNet ++由编码器和解码器组成,它们通过一系列嵌套的密集卷积块连接。 UNet++的主要想法是在融合之前在编码器和解码器的特征图之间架起语义鸿沟。 例如,(X0,0,X1,3)之间的语义鸿沟是使用具有三个卷积层的密集卷积块来架起的。 图中,黑色表示原始的U-Net,绿色和蓝色表示跳跃路径上的密集卷积块,而红色表示深监督。 红色,绿色和蓝色组件将UNet ++与U-Net区别开来。 (b)详细分析UNet ++的第一个跳跃途径。(c)如果在深监督下进行训练,则可以在预测时对UNet ++进行剪枝。
提出的网络架构:UNet++
图1a显示了所建议架构的高级概述。 如图所示,UNet ++从编码器子网络或骨干网开始,然后是解码器子网络。 UNet ++与U-Net(图1a中的黑色部分)的不同之处在于重新设计的跳跃路径(以绿色和蓝色显示)连接了两个子网络,并使用了深度监督(红色)。
重新设计的跳跃路径
重新设计的跳跃路径可改变编码器和解码器子网络的连通性。 在U-Net中,编码器的特征图直接在解码器中被获取; 但是,在UNet ++中,它们经历了密集的卷积块,其卷积层数取决于金字塔级别。 例如,节点X0,0和X1,3之间的跳跃路径由一个具有三个卷积层的密集卷积块组成,其中每个卷积层之前都有一个串联层,该层将同一密集块的先前卷积层的输出与较低层密集块的相应上采样输出融合。 本质上,密集卷积块使编码器特征图的语义级别更接近解码器中等待的特征图的语义级别。 假设是,当接收到的编码器特征图和相应的解码器特征图在语义上相似时,优化器将面临更容易的优化问题。
形式上,我们将跳跃路径的公式表示如下:令xi,j表示节点Xi,j的输出,其中i沿编码器索引下采样层,j沿跳跃路径索引密集块的卷积层。 xi,j表示的特征图的堆栈计算如下
其中函数H(·)是卷积运算,后跟激活函数,U(·)表示上采样层,[]表示concat。基本上,j = 0的节点仅接收来自编码器上一层的一个输入;级别j = 1的节点从编码器子网络接收两个输入,但在两个连续级别上;并且级别j> 1的节点接收j + 1输入,其中j输入是同一跳跃路径中前j个节点的输出,最后一个输入是来自较低跳跃路径的上采样输出。所有先前的特征图都会累积并到达当前节点的原因是因为我们沿每个跳跃路径使用了密集的卷积块。 图1b进一步阐明了等式,通过显示特征图如何通过UNet ++的顶部跳跃路径来实现图1。
深监督
我们建议在UNet ++中使用深监督[6],使模型可以在两种模式下运行:1)精确模式,其中平均所有分割分支的输出; 2)快速模式,其中最终分割图仅从分割分支之一中选择,如何选择确定了模型剪枝的程度和速度增益。 图1c显示了在快速模式下选择分割分支如何导致复杂性变化的网络架构。
由于嵌套的跳跃路径,UNet ++会在多个语义级别{x0,j,j∈{1,2,3,4}}上生成全分辨率的特征图,这些图可以进行深监督。 我们在上述四个语义级别中的每个语义级别上都添加了二进制交叉熵和dice系数的组合作为损失函数,其描述为:
其中Yˆb和Yb分别表示b^th图像的预测概率和ground-truth, N表示batch大小。
总而言之,如图1a所示,UNet ++与原始U-Net在三个方面有所不同:1)在跳跃路径上具有卷积层(以绿色显示),这架起了编码器和解码器特征图之间的语义鸿沟; 2)在跳跃路径上具有密集的跳跃连接(以蓝色显示),从而改善了梯度传递; 3)具有深监督(以红色显示),这将在第4节中显示,可以进行模型剪枝并进行改进,或者在最坏的情况下,可以实现与仅使用一个损失层相当的性能。
实验
数据集
如表1所示,我们使用四个医学成像数据集进行模型评估,涵盖了来自不同医学成像模式的病变/器官。 有关数据集和相应数据预处理的更多详细信息,请读者参看补充材料。
基线模型
为了进行比较,我们使用了原始的U-Net和自定义的宽U-Net架构。 我们选择U-Net,因为它是图像分割的通用的性能基准。 我们还设计了一个宽U-Net,其参数数量与我们提出的架构相似。 这是为了确保我们的架构所获得的性能提升不仅仅是因为参数数量的增加。 表2详细列出了U-Net和宽U-Net体系结构。
实施细节
我们评估dice系数和IoU,并在验证集上使用了早停法。 我们还使用了学习率为3e-4的Adam优化器。 表2中显示了U-Net和宽U-Net的架构详细信息。UNet++是从原始U-Net架构中构建的。 沿跳跃路径(Xi, j)的所有卷积层都使用大小为3×3(对于3D肺结节分割,大小为3×3×3)的k个卷积核,其中k = 32×2^i。 为了进行深监督,将1×1卷积层和sigmoid激活函数附加到每个目标节点:{x0, j | j∈{1,2,3,4}}。 结果显示,UNet ++在给定输入图像的情况下生成四个分割图,并将其进一步平均以生成最终的分割图。 可以在github.com/Nested-UNet上找到更多详细信息。
结果
表3在数量参数和分割精度方面对U-Net,宽U-Net和UNet ++进行了比较,以完成肺结节分割,结肠息肉分割,肝脏分割和细胞核分割的任务。如图所示,宽U-Net始终胜过U-Net,除了肝脏分割(两种结构的性能相当)之外。这种改进归因于宽U-Net中大量的参数。没有深监督的UNet ++可以在U-Net和宽U-Net上实现显著的性能提升,平均IoU提升2.8和3.3点。具有深监督的UNet ++与没有深监督的UNet ++相比平均提高了0.6点。具体而言,使用深监督可显著改进肝和肺结节的分割,但是这种改进对于细胞核和结肠息肉分割消失了,这是因为息肉和肝脏在视频帧和CT切片中出现的尺度不同。因此,使用所有分割分支(深监督)的多尺度方法对于准确分割至关重要。图2显示了U-Net,宽U-Net和UNet ++的结果之间的定性比较。
模型修剪
图3显示了采用不同级别的剪枝后的UNet ++的分割性能。 我们使用UNet ++ L^i表示在i级修剪的UNet ++(有关更多详细信息,请参见图1c)。 如图所示,UNet ++ L^3的推理时间平均减少了32.2%,而IoU仅降低了0.6点。 更积极的修剪进一步减少了推理时间,但代价是精度大大降低。
总结
为了满足对更精确的医学图像分割的需求,我们提出了UNet ++。 建议的架构利用了重新设计的跳跃路径和深监督的优势。 重新设计的跳跃路径旨在减小编码器和解码器子网络的特征图之间的语义鸿沟,从而使优化器面临更简单的优化问题。 深监督还可以实现更精确的分割,尤其是对于出现在多个尺度的病变(例如结肠镜检查视频中的息肉)。 我们使用涵盖肺结节分割,结肠息肉分割,细胞核分割和肝分割的四个医学影像数据集评估了UNet ++。 我们的实验表明,经过深监督的UNet ++分别比U-Net和宽U-Net分别获得3.9和3.4点的平均IoU增益。
参考文献
-
S. G. Armato, G. McLennan, L. Bidaut, M. F. McNitt-Gray, C. R. Meyer, A. P. Reeves, B. Zhao, D. R. Aberle, C. I. Henschke, E. A. Hoffman, et al. The lung image database consortium (lidc) and image database resource initiative (idri): a completed reference database of lung nodules on ct scans. Medical physics, 38(2):915–931, 2011.
-
M. Drozdzal, E. Vorontsov, G. Chartrand, S. Kadoury, and C. Pal. The importance of skip connections in biomedical image segmentation. In Deep Learning and Data Labeling for Medical Applications, pages 179–187. Springer, 2016.
-
D. Fourure, R. Emonet, E. Fromont, D. Muselet, A. Tremeau, and C. Wolf. Residual conv-deconv grid network for semantic segmentation. arXiv preprint arXiv:1707.07958, 2017.
-
K. He, G. Gkioxari, P. Dolla ́r, and R. Girshick. Mask r-cnn. In Computer Vision (ICCV), 2017 IEEE International Conference on, pages 2980–2988. IEEE, 2017.
-
G. Huang, Z. Liu, K. Q. Weinberger, and L. van der Maaten. Densely connected convolutional networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, volume 1, page 3, 2017.
-
C.-Y. Lee, S. Xie, P. Gallagher, Z. Zhang, and Z. Tu. Deeply-supervised nets. In Artificial Intelligence and Statistics, pages 562–570, 2015.
-
X. Li, H. Chen, X. Qi, Q. Dou, C.-W. Fu, and P. A. Heng. H-denseunet: Hybrid densely connected unet for liver and liver tumor segmentation from ct volumes. arXiv preprint arXiv:1709.07330, 2017.
-
J. Long, E. Shelhamer, and T. Darrell. Fully convolutional networks for seman- tic segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 3431–3440, 2015.
-
O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234–241. Springer, 2015.
-
N. Tajbakhsh, J. Y. Shin, S. R. Gurudu, R. T. Hurst, C. B. Kendall, M. B. Gotway, and J. Liang. Convolutional neural networks for medical image analysis: Full training or fine tuning? IEEE transactions on medical imaging, 35(5):1299–1312, 2016.
-
Z. Zhou, J. Shin, L. Zhang, S. Gurudu, M. Gotway, and J. Liang. Fine-tuning convolutional neural networks for biomedical image analysis: actively and incre- mentally. In IEEE conference on computer vision and pattern recognition (CVPR), pages 7340–7351, 2017.