DoubleU-Net: A Deep Convolutional Neural Network for Medical Image Segmentation
用于医学图像分割的Double U-Net框架
Abstract
语义分割就是为图像中每一个像素点的类别进行预测。基于编解码结构的分割网络,如UNet及其变体,是医学图像分割中广泛使用的网络。为了进一步提升UNet在不同分割任务中的性能,本文提出了一种新的Double-UNet框架,即两个UNet网络结构叠加的组合。
其中第一个UNet使用预训练的VGG-19作为编码网络,可以将从ImageNet学习到的特征迁移到其他任务中;为了捕获更多的语义信息本文又加入了另一个UNet网络,使用ASPP空间金字塔池化来捕获上下文信息。
本文将Double-UNet在四类医学分割数据集上进行了测试,测试图像包含多种模态:结肠镜、皮肤镜、显微镜图像。实验结果显示Double-UNet的分割性能均优于原始的UNet结构,尤其在MICCAI2015等挑战赛中对于微小扁平的息肉分割都具有更精确的分割结果,在UNet基础上有较大改进。基于Double-UNet在不同医学图像分割任务中均取得了令人振奋的效果,显示出Double-UNet可以作为医学图像分割和跨数据集评估的基准模型,用以衡量深度学习模型的泛化性。
Section I Introduction
医学图像分割负责将医学图像RoI区域中的每个像素进行标记,这通常是进行临床诊断的关键任务,不管是进行损伤、病变的计算机自动诊断亦或是制定治疗计划。医学图像分割可以帮助医生聚焦于病变的特定求、提取用于诊断的细节信息。医学图像分割的主要挑战在于难以获取大量的标记数据、缺少高质量的标注信息用于训练,以及缺少分割的评判依据、而不同病人的图像差异大等。因此需要建立一种自动、通用、高效的语义分割方法。
卷积神经网络已经在医学图像分割诸多任务中达到了SOTA。一种最常用于语义分割方法任务的框架是全卷积神经网络(FCN),UNet是另一种支持端到端训练的分割框架,包含两部分:分析路径和合成路径。analysis path负责学习图像特征,synthesis path则根据学习到的特征生成分割图谱。除此之外UNet还引入了跨层连接skip connection从而使得再analysis path中学习的信息可以传递到对应的synthesis path从而可以产生更加精确的分割结果。因此,通过增加UNet的层数网络可以学习到更具表征性的特征,从而产生精确的分割掩膜。
泛化性指的是模型在独立数据集上的性能;鲁棒性则是模型在具有挑战性测试图像上的性能,这两者均是应用于临床的AI系统的重要指标。因此建立一个强有力的医学图像框架具有足够的泛化性和鲁棒性,从而适应不同的医学图像应用。在ImageNet上预训练过的模型可显著提升性能,VGG-19就是其中之一,除此之外受UNet及其变体在医学图像分割的启发,本文提出使用VGG-19和改进后的UNet作为网络的编码部分,因为涉及两个网络,本文称之为DoubleU-Net。
选用VGG-19的原因有:
(1)与其他预训练网络相比,VGG-19更加轻量级;
(2)VGG-19的结构与UNET相近,方便与UNet进行级联;
(3)通过搭建更深层次的网络产生更加精确的分割结果。
本文的主要工作总结如下:
(1)提出了一种用于语义分割的新的框架-DoubleUNet,依次通过两个UNet的编解码结构,第一个则是在ImageNet上预训练的VGG-19网络,并且使用了空间金字塔池化(ASPP),网络的其余部分则是从头一点一点构建的。
(2)本文框架在4种不同医学图像数据集上的进行了实验,分别是2种结肠镜、皮肤镜、显微镜图像。
(3)在4种不同数据集上的实验结果显示DoubleUNet与UNet相比性能显著提升,可以作为医学图像分割模型中新的基准模型。
Section II Related Work
在不同的CNN框架中,编解码网络-FCN及其拓展模型UNet广泛用于2D图像的语义分割,后续也有研究人员通过空洞卷积聚合多尺度的上下文信息进一步提升分割精度。
DeepLab也是一类语义分割框架,其中DeepLabV3的synthesis path就包含了更少的卷局曾,但也是用了与UNet类似的skip connection。
还有借助global pyramidal features设计的稀疏网络用于复杂场景理解,自己DeepResidual UNet\DRINET等诸多变体。
在ResUNet++中,作为ResUNet的增强版本,在网络中加入了SE Block和ASPP、注意力模块,从而提升分割精度。从上述相关工作可以看著,在自然图像分割和医学图像分割方面均在大力发展不同的网络框架。
近期越来越多的工作聚焦于设计更具泛化性的分割网络,因此越来越多的工作在不同的自然图像和医学图像数据集上进行测试。但是AI在医学领域的应用仍处于起步阶段,一大限制就是缺少大型的数据集,并且目前能获取的数据集大多类别不均衡。
也许能在自然图像取得可以接受的性能,但是有很多十分具有挑战性的医学图像,比如结肠镜中的扁平息肉通常就在检查中被遗漏,如果早期遗漏有可能发展为癌症。需要设计更加精确的分割网络来分割这些具有挑战性的医学图像,因此本文提出了DoubleUNet。
Section III Double UNet
Architecture
Fig 1展示了DoubleUNet的整体结构,首先是NETWORK1:VGG-19的编码网络,蓝色则是ASPP模块,随后经过解码网络,并且使用了SE Block。NETWORK1的输入和输出还进行了element-wise multiplication。
而NETWORK2与UNet的区别仅在于使用了ASPP和SE Block。
在NETWORK1中输入和输出(OUTPUT1)相乘后的结果作为NETWORK2的输入,随后将两个网络的输出(OUTPUT1和OUTPUT2)进行级联(concatenate),这样可以查看OUTPUT1与最终分割结果的差异。
因为本文认为NETWORK1输出的feature map依旧可以通过再次获取原始输入图像和对应的mask进一步进行改进。这也是本文使用两个UNet得到动机。而SE Block的作用则是减少冗余信息,只传递最相关的信息。而采用ASPP则是因为其可以有效地提取高分辨率的feature map从而提升分割精度。
Part A Encoder Explanation
NETWORK1使用VGG-19作为编码网络,在NETWORK2中则是一级一级搭建成的编码网络,每一层包含两次3x3卷积->BN->Relu。随后会经过SE Block用于提升feature map的质量,随后通过maxpooling降低空间维度编码结构负责对输入图像中的信息进行编码。
Part B Decoder Explanation
两个网络的Decoder部分则和原始UNet没有过多改动,每一个解码模块会经过2x2的双线性插值用于恢复空间分辨率,随后在通过skip connection将同层次的解码网络链接过来。细节在于NETWORK1的skip connection只连接了自己的decoder而NERWORK2则连接了两个编码网络的信息。插值之后依旧经过两次3x3卷积以及BN和ReLU激活以及SE Blcok。最后通过卷积和sigmoid激活获得最终的segmentation mask。
Section IV Experiments
Part A Datasets
本文在四类数据集上对DoubleUNet的性能进行了测试,分别是:
MICCAI 2015挑战赛:
training:CVC-ClincDB
testing:ETIS-Larib
ISIC-2018挑战赛-LeisionBoundary Segmentation
2018 Data Science Bowl-nuclei segmentation
数据集划分:train:val:test = 80%:10%:10%
详细信息参见TableI
Part B Evaluation Matrics
DSC,mIoU,Precision,Recall
Part C Experiment setup and configuration
框架:Keras TF2.1.0
加速卡:Vlota 100 GPU
损失函数:交叉熵损失函数
epoch:300
早停法+ReduceLROnPlateau
Part D Data Augmentation
鉴于医学图像数据的稀缺性,常使用数据增强来增加训练时的样本数目。本文则使用了center crop,random rotation,transpose,elastic transform等数据增强手段,因此一张原始训练数据扩增为25张不同的数据。
Section V Results
本文依旧与原始的UNet基准模型进行对比,采用了同样的实验设置。
Part A MICCAI 2015
Table II展示了在MICCAI2015分割挑战赛中的分割结果,可以看到DoubleUNet比原始UNet分别提升了:DSC 6.07%,mIoU:1.31%;并且也能看出在ImageNet上预训练过的模型与从零开始训练的网络相比,具有更好的跨数据集性能。
根据Fig 2的可视化结果也能看出最终的分割结果OUTPUT2比中间结果OUTPUT1更好。
Part B CVC-Clinic DB
Table III展示了在CVC-Clinic DB的分割结果,依旧比其他SOTA网络有提升,并且通过Fig 3的可视化结果可以看出DoubleUNet在一些具有挑战性的图像(如息肉)具有更好地分割结果,而这一部分经常在结肠镜检查中遗漏。
Part C Lesion Boundary Segmentation
Table IV展示了病灶边缘分割的结果,根据Fig 4的可视化结果可以看到不管是DoubleUNet的中间结果还是最终结果都比比UNet有显著提升,但还是OUTPUT2的分割精度更高。
Part D 2018 Data Science Bowl challenge dataset
结果参见TableV和Fig5,与之对比的则是UNet++框架。
Section VI Discussion
Table VI展示了DoubleUnet 与UNet的性能对比,在所有测试的数据集上DoubleUnet都较UNet有提升,主要可以得到以下观察结果:
(1)MICCAI 2015提供了跨数据集泛化性的对比,DoubleUNet泛化性更佳;
(2)ImageNet预训练后的网络比从头训练的网络性能更佳;
(3)DoublUNet在挑战性图像上具有更好的分割性能。
总的来说,DoubleNet对于不同数据集、多器官或者多中心的医学图像具有更好的分割性能,充分显示了模型的鲁棒性;此外通过预训练网络进行迁移学习也是十分有意义的。
Section VII Conclusion
本文提出了DoubleUnet用于医学图像分割,主要包含5部分:两个UNet网络,VGG-19,SE Blcok,ASPP.实验结果表明DoubleUNet性能明显优于UNet.
并且这一结构更加灵活多变,可以方便的将其他模块集成到DoubleUNet中,本文相信这样可以进一步提升DoubleUNet的分割精度。
未来本文计划建立一个适用于不同医学图像的分割模型,在保证精度的前提下尽可能精简结构。因为DoubleUnet的局限之一就是引入了过多的网络参数,这样会增加训练时间。因此本文未来会进一步精简网络结构。