医学图像分割新网络：Boundary-aware Context Neural Network for Medical Image Segmentation

作者团队：中国科学院&中国科学技术大学

摘要

BA_Net：医学图像语义分割网络：在多个数据集上，表现SOTA！医学图像分割对于进一步的医学图像分析和疾病诊断提供了一个可信赖的基础。近年来，随着卷积神经网络的发展，医学图像分割取得了很大的进展，但是，基于CNN的方法在目标的边界处并没有达到理想的分割 效果。这是由于上下文信息的局限以及经过卷积和池化之后的特征图不具有充足的判别力所造成的的。医疗图像的特点是具有很高的类内差异性，以及类间的区分不明及噪声。因此如何提取有效的上下文信息以及如何聚合区分不明的特征对于精细的分割仍是一个很大的挑战。
在这篇文章中，作者提出了BA-Net，一种用于2D医疗图像分割的方法，捕捉了更丰富的上下文信息，提供了更有效的空间信息。BANet采用了encoder-decoder的结构。在每个encoder阶段，应用了金字塔边缘提取模块。然后作者设计了一个mini的多任务学习模块用于联合学习分割目标的mask和检测病变边缘。特别的是，为联合两个任务。提出了一种新的交互式注意力模块，有效利用了边界信息，为分割预测提供了有力的线索信息。最后，提出了一个特征聚合模块，用于选择性的聚合encoder模块的多级特征。通过级联以上三个模块，在每个阶段中都编码了足够的上下文信息和精细化的特征。实验结果表明，该方法在五个数据集中都达到了SOTA。

一、引言及相关方法

图像分割在医疗图像分析中具有重要的作用，旨在对每个像素进行分类，并得到精细化的病变区域识别。近年来，随着医疗图像技术和设备的发展和普及，超声、MRI、CT及其他设备对于疾病的诊断提供了有效的手段。这些技术已经广泛应用到日常的医疗研究和诊断治疗中。对于不同种类的医疗设备，图像是一种重要的图像分析手段，比如肺部图像分割。精确的病检测对于医学诊断提供了重要且可靠的参考价值。
随着医学图像数量的增加以及AI技术的发展，计算机辅助诊断成为一种重要的手段，然而自动的病变（器官或是组织）识别仍是一个复杂且极富挑战性的工作。
首先，病变区域具有不同的尺寸和形状。对于某些疾病，不同的个体差异增加了识别的难度。

上图所示分别为皮肤病变和结肠息肉。
其次是，对于病变组织和背景之间的低区分性也带来了很大的挑战，尤其是，需要关注的区域同时包含着复杂的组织和器官，这对于区分边界像素带来了极大的困难。除此之外，某些设备和图像的噪声对于精确的分割也具有一定的阻碍。
在过去的数十年间，提出了很多的医疗图像分割的自动处理分析算法，主要可以分为三类：

• 基于灰度图像的；
• 基于纹理的；
• 基于atlas-based，翻译为地图集，也就是人工标记的图像。
  虽然这些方法已经取得了很好的效果，但仍有以下几点不足：
• 传统方法通常会利用低维的手工设计的特征，通常会限制复杂场景的预测表现，同时，大量的可利用的原始图像信息被忽略了；
• 鲁棒性低，依赖于有效的数据预处理。
  近年来，随着卷积神经网络的发展，大多数医学图像分割方法采用encoder-decoder的网络结构，其中，最具代表性的即UNet和FCN。为了有效的识别目标，提取及聚合高维语义特征和低维特征是很有必要的。
  受以上分析的启发，本文提出了一种新颖的基于边界已知BA-Net的医学图像分割方法，

二、本文主要方法

2.1 总体介绍

如下图所示，BA-Net采用encoder-decoder的网络结构，并且在backbone阶段采用Resnet（使用ImageNet进行预训练）。在encoder结构中，ResNet中的最后一个全局池化层及全连接层被去掉了，近保留一个卷积和四个残差模块用于主要的特征提取作用。不失普遍性的，对于每张输入图像，将四个残差block的输出分别表示为$F_i$, $i\in 1,2,3,4$.
为了增加感受野，最后两个残差模块中使用的空洞卷积，空洞系数为2，并且保留与前一层模块相同的空间像素，因此，每个模块的输出尺寸为输入尺寸的$1/4,1/8,1/8,1/8$。此外，在最后一个残差模块的上层特征图中应用了ASPP即空洞空间金字塔池化模块用于编码和捕捉多尺度特征。ASPP模块包含四个并行的具有不同空洞系数的空洞卷积和一个全局平均池化，主要用于压缩和合并特征图。为了产生更丰富的上下文信息，用于解码过程的引导，本文增加了三个模块用于挖掘同一水平的特征以及聚合不同水平的特征在每个encoder阶段。

• PEE（Pyramid edge extraction）金字塔边缘提取模块用于聚合不同尺寸的边缘信息；
• MTL（mini Multi-task learning module）mini多任务学习模块，通过利用相关边缘检测和分割任务中的潜在的关联信息和辅助性特征来得到更加丰富的知识信息。
• CFF（Cross feature fusion）交叉特征融合模块，聚合不同level之间的特征。
  最后，在decoder部分，通过聚合ASPP模块的不同输出$D_i,i\in 1,2,3,4$，然后依次对每个阶段的编码特征进行最后的分割预测。
  .

2.2 Pyramid Edge Extraction Module

病变区域的边缘对于分割目标的位置提供了重要的信息。然而，该边缘也是非常复杂和多变的，为了得到一个鲁棒的边缘信息，本文设计了一个简单却有效的特征提取组合，能够挖掘不同尺寸的边缘特征。如上图所示，首先使用一个$1X1$的卷积压缩backbone中每个阶段的最后一个残差模块，并将他们作为PEE模块的输入，定义如下：

其中，$F_i^{\prime }$表示每个残差模块的缩减后的特征图，函数表示$1X1$的卷积运算，${\theta }_i$表示分别的参数，通过减去不同尺寸的的平均池化后的value值，得到不同尺寸粒度的边缘特征，不失一般性的，定义$S$为池化操作：

其中，$F_{i,p}^{(s)}$表示当前第$i$阶段的使用$s$个池化操作的边缘特征，并且，$av{g}_s$表示相关平均池化操作。为了整合得到的金字塔边缘特征，本文通过concat操作进行特征聚合，最后使用$1X1$的卷积进行组合。

其中，$C$表示concat过程，$F_{i,p}$表示PEE模块的输出特征图，${\theta }_{i,p}$表示相关参数。通过设计这样一个具有多尺度的特征聚合模块，可以增强对应level的表示能力。

2.3 Mini Multi-Task Learning Module

来自于目标边缘的额外信息可以帮助判断目标的形状，并且，语义分割和边缘检测具有很强的依赖关系。基于这个观点，本文提出一个mini的多任务学习网络mini-MTL，将其嵌入在每个阶段中，这样可以不需要引入过多的参数即可进行分割和边缘检测。这个mini-MTL模块的主要作用是通过潜在的关联来得到一定的效果增益。如上图所示，这个多任务网络包含两个主要组成部分：the task specific branch 和 interactiv attention layer。每个branch包含两个卷积层和一个上采样层。卷积操作主要用来编码任务相关的特征，上采样层是用来得到对应预测的mask。在第$i$阶段，PEE模块的特征图$F_{i,p}$是作为两个子任务的输入，用于提取同时相关任务的特征。

其中$F_{i,ϵ}^l$和$F_{i,S}^l$表示边缘提取子任务和分割子任务中提取的第$l$层特征图，$l\in 1,2$。另外，交互式的注意力模块(IA)被设计在第一层卷积中，用于挖掘不同任务之间的交互信息。
如下图所示，为了整合其他任务的有效信息，本文设计了一个简单却有效的交互式注意力整合方法。以边缘特征整合为例，首先利用一个sigmoid函数得到一个权重mask，即表示当前边缘特征的重要位置。然后对注意力权重取反，即得到另一个任务的权重注意力mask。最后，我们可以从分割特征中选择性的发送有用的信息到当前的边缘特征中，通过逐像素的相乘操作。

这里，边缘特征图和分割特征图均以二值化的形式表示，loss函数均采用二值交叉熵损失函数：

以上分别表示边缘检测和分割两个子任务的损失函数。

其中，$P_j$和$G_j$表示第$j$个像素的预测边缘特征图和Ground truth。$N$表示像素的数量，因此总体loss可表示如下：

2.4 Cross Features Fusion Module

在encoder部分，低层特征具有蜂丰富的空间信息，高层特征具有丰富的语义信息。为了能够同时利用空间结构的细节特征和语义特征，本文提出了一个CFF（cross feature fusion module），能够选择性的聚合高层和低层特征。如下图所示：

其中，对于第$i$个特征图$F_{i,M}$，CFF模块自适应的从多个输入特征中选择互补的内容，通过以下注意力机制：

$\sigma$表示sigmoid激活函数，$⨂$表示逐像素相乘，因此，来自于不同level的特征通过CFF模块整合起来，而且能够避免引入过多的信息。

2.5 Decoding and Optimization

通过在encoder阶段级联以上三个模块，丰富的上下文信息和精细化的特征被编码进来。在decoder网络中，通过聚合ASPP模块的输出特征和每个阶段在的编码特征得到解码特征$D_i,i\in 1,2,3,4$:

其中，$D_i$表示每个阶段的编码特征，$F$表示$1X1$卷积，${\theta }_{i,D}$表示相应的参数。
在整个网络中，利用的标准的二值交叉熵损失函数来最小化误差。在端到端的训练过程中，完整的损失函数定义如下：

其中，$L_D$表示decoder部分的损失，${\lambda }_i$表示平衡系数。

三、实验

3.1 医学图像数据集介绍

为验证本文方法，在五个医学图像数据集中进行了实验验证，包括皮肤镜像图、内窥镜图、X射线图和视网膜眼底图像。包括以下五个数据集：

• ISIC-2017,皮肤镜像图像，2000张训练集图像，600张测试集图像，图像尺寸为540x722到4499x6748不等；
• Kvasir-SEG，1000张息肉图像；
• CVC-ColonDB，包含380张结肠镜图像，尺寸为574x500;
• SZ -CXR ,包含566张X射线图像；
• RIM-ONE-R1，包含169张视网膜眼底图像。

3.2 参考模型

本文中，与当前六个SOTA模型进行了实验对比，包括FCN、U_net、MultiResUNet、AG-net、CE-Net以及DeeplabV3. 为进行公平的实验对比， 对比方法采用作者开源代码。所有的实验采用相同的数据预处理，并且使用相同的评价指标。

3.3 评价指标

评价指标包括 Dice Similarity Coefficient (DI), Jaccard Index(JA), Accuracy (AC), ensitivity (SE) and Specificity (SP)。

3.4 与SOTA的实验对比

通过与上述多个SOTA医学图像分割方法在五个数据集进行实验对比，下表表明，本文方法在不同数据集上均达到很好的分割效果。

下图所示为本文方法与其他SOTA的医学图像分割方法的分割可视化效果图，从图中可以看出，本文方法相对来说，分割精度更高，尤其是在边缘处的分割更加精确。

四、结论

本文提出了BA-Net网络，用于医学图像分割，通过本文方法，提取了更加丰富的上下文信息，网络结构中使用了三个级联的新颖模块，在五个数据集上均达到了SOTA。