基于注意力机制的U-Net脑脊液细胞分割

代 茵, 刘维宾, 董昕阳, 宋雨朦

(1.东北大学 医学与生物信息工程学院, 辽宁 沈阳 110169; 2.东北大学 教育部医学影像与智能分析工程中心, 辽宁 沈阳 110169; 3.约克大学 计算机学院, 英国 约克郡 YO10 5DD)

目前国内外针对自身免疫性脑炎的诊断研究是利用磁共振成像(magnetic resonance imaging, MRI)检测以及结合其他影像学诊断方法，如正电子发射计算机断层显像(positron emission computed tomography, PET)、超声波检查，以及脑脊液(cerebrospinal fluid, CSF)检测，而CSF检测是最为标准的，通常作为医生判断疾病的“金标准”[1].免疫机制在精神疾病病理学中的作用越来越得到人们的关注，多项研究表明了自身免疫性疾病与精神分裂症之间存在着必然的联系.CSF分析对识别中枢神经系统是否存在炎症起着核心的作用[2].CSF检测结果不仅可以作为判断脑炎与精神分裂症的依据，还可以作为临床孤立综合征[3]、额颞痴呆、小儿鞍上生殖细胞瘤、多发脑梗死性痴呆、神经莱姆病、遗传性多发脑梗死性痴呆、单发脑梗死性痴呆等疾病的重要依据.

CSF细胞图像有其独有的特点，其存在细胞边界模糊与背景不易区分、数据图像存在一定的噪声、相同类型的细胞通常在图像中成簇存在[4]、样本数据过少等问题.目前医生大多还是通过手动分割CSF细胞的方式并通过经验来诊断相关疾病，工作强度大，效率低，而基于人工智能的传统分割算法的鲁棒性差，所以针对CSF细胞检测的自动分割研究显得尤为重要.近几年深度学习发展迅猛，有研究人员在图像分割与深度学习相结合的领域作了许多研究与贡献.在CSF图像的研究中，Xu等[4]提出一种层次非最大抑制(hierarchy-non maximum suppression,Hierarchy-NMS)算法，与非极大抑制算法和软化非极大抑制算法进行了比较，证明了层次非最大抑制算法在细胞识别和计数方面均优于非极大抑制的一些变体.但该研究更多的关注点是识别与记数，而没有考虑CSF细胞边界与背景精确分割的问题.

在医学图像领域内，Jiang等[5]设计了一种基于双通道双向LSTM(dual channel bidirectional LSTM, DC-BLSTM)的神经网络，将3D图像分割任务转化为多个1D分割任务，比现有的基于卷积神经网络的3D神经元图像分割方法更容易标记训练样本.另外极具代表性的是以U-Net模型为基础框架，通过改造编码器或解码器实现其性能的提升.研究人员通过这些变体对许多细胞图像进行了分割研究，例如对腺体细胞进行分割[6]，肌肉干细胞的分割[7]，肺、细胞轮廓、胰腺的分割研究[8]，多器官分割和心脏分割[9]，以及对ISBI细胞追踪挑战赛2019数据集的分割测试[10]等.注意力机制的优秀性能与良好表现引起了研究人员的注意.Vaswani等[11]提出了Transformer网络架构，此架构是完全基于注意力机制，并且在WMT 2014 English-German数据集上验证了优越性能.Zhang等[12]提出一种带有注意力模块的O形网络结构，可以在不分割的情况下检测生物医学图像中的节点，使用该方法在两个视网膜数据集和一个神经元数据集中取得了良好的检测效果.迁移学习可以提升模型泛化能力，使模型的初始性能得到提高.刘珍梅[13]将迁移学习算法与宫颈细胞核分割算法相结合，减少分割标注数据集制作的压力，得到的模型收敛性更好.

本文将注意力机制与深度学习相结合，并引入迁移学习，对CSF细胞进行精确分割，取得了理想的结果.

1 方法描述

本文采用了基于注意力机制的U-Net卷积神经网络对CSF细胞进行精确分割，并引入预训练模型进行迁移学习[14]，在一定程度上提高了模型训练速度与性能.概括来说实验流程共分为4个部分：

1) 数据集扩充：将从已有的CSF图像中筛查出29例可靠图像，然后再采用对其旋转、裁剪、镜像的方法进行扩充；

2) 数据预处理：将扩充后的数据集进行尺寸大小统一预处理，使其适配batch的读取方式，网络的感受野也就限定了统一的尺度范围，方便统一做训练，并且训练集通过拉普拉斯算子进行了滤波处理；

3) 网络训练：将处理好的数据集传入带有注意力机制的U-Net卷积神经网络模型中，并引入迁移学习进行训练；

4) 分割指标评估：训练完成后，使用训练模型对数据进行预测，使其预测图与Ground Truth进行指标评估，评估指标包括相似系数值、平均交并比和类别平均像素准确率.

实验流程示意图如图1所示.

图1 实验流程图

2 基于注意力机制的U-Net网络 结构模型

2.1 总体设计

本文所使用的基于注意力机制的U-Net卷积神经网络总体分为编码器和解码器，在解码器中加入注意力模块，在对编码器每个分辨率上的特征与解码器中对应特征进行拼接之前，使用Attention Gate重新调整了编码器的输出特征，该模块生成一个门控信号，用来控制不同空间位置处特征的重要性[15].网络的总体结构如图2所示.

图2 基于注意力机制的U-Net

2.2 U-Net网络结构

2.2.1 编码器结构

作为主干特征提取网络，部分是由卷积层和最大池化层组合而成的模块，整体结构与VGG[16]极为相似，所以本文采用VGG16的主干特征提取网络，如此也方便使用预训练权重进行迁移学习.

在编码器中，共进行了5次卷积模块的操作，输入图像大小为512×512×3，前2个卷积模块操作分别为2次卷积核为3×3的64通道和128通道的卷积操作，获得的初步有效特征层再分别进行一次2×2的最大池化，获得相应的特征层.后3个卷积模板操作分别为3次卷积核为3×3的256通道、512通道和512通道的卷积操作用来获得初步有效特征层.第3个卷积模块与第4个卷积模块再分别进行一次2×2的最大池化，第5个卷积模板不再做最大池化操作.编码器的结构图如图3所示.

图3 编码器的结构图

2.2.2 解码器结构

经过编码器的主干特征提取获得5个有效特征层，对这5个特征层进行特征融合，所采用的融合方式是对特征层进行上采样并堆叠.为了方便搭建网络以及考虑到通用性，本文设计的U-Net与Ronneberger等所提出的U-Net[17]略有区别.本文采用上采样时先进行2倍的上采样，然后进行特征融合，这样做的优势是最终得到的特征层与输入图片大小相同.解码器的结构图如图4所示.

图4 解码器的结构图

2.3 注意力机制

在传统的U-Net中，为了避免在解码过程中丢失重要细节信息，使用了跳跃连接的方式，将编码过程中提取得到的映射直接链接到解码器相对应的层.但是这样操作有个缺点，即提取的低级特征会存在很多冗余信息.利用注意力机制[18]能够抑制模型学习与任务不相关的部分，强化学习与任务相关的特征，从原理上分为空间注意力与通道注意力.

Jaderberg等[19]提出的STN网络可以根据任务的不同，自适应地将数据进行空间变换.此网络是基于空间的注意力思想.SENet[20]在ILSVRC 2017分类中，赢得了第一名，此网络本质上是基于通道的注意力模型，它只使用全局平均池化来压缩通道特征.通过学习的方式自动获取每个特征通道的重要性，然后通过获取到的通道重要程度来提升有用信息和抑制无用信息.

本文采用的是通道注意力模块与空间注意力模块相结合的方式，将模块添加到网络的解码器中，通道的每一个高级特征都可以看作是一种特定的响应．通过挖掘通道映射之间的相互依赖关系，可以强调相互依赖特征映射，改善特定语义的特征表示．通道注意力模块与空间注意力模块相结合的模块示意图如图5所示[15].

图5 注意力模块

通道注意力专注于有意义的输入图像，为了有效地计算通道注意力，需要压缩输入特征图的空间维度.对于空间信息的聚合，常使用的方法是平均池化，而最大池化可以收集到难以区别物体之间的重要特征线索，来获取更加详细的通道注意力，所以通道注意力模块同时使用平均池化和最大池化.然后将它们送入一个权重共享的多层感知机中，最后再将这些输出的特征按对应位置相加.通道注意力模块如图6所示[15].

图6 通道注意力模块

空间注意力专注于输入图像的有效信息的位置，可以很好地弥补通道注意力的不足.为计算空间注意力，沿着通道轴使用平均池化和最大池化，将它们相连，生成一个有效特征描述符.最后通过卷积层生成空间注意力.空间注意力模块如图7所示[15].

图7 空间注意力模块

2.4 损失函数

本文采用交叉熵与Dice损失相结合的形式作为语义分割的损失指标.交叉熵损失函数常用于分类问题，本文是对图像细胞做分割，本质也是对细胞与背景像素点做二分类，在二分类情况下，模型最后需要预测的结果有两种情况，对于每个类别预测出的概率为pi和1-pi，交叉熵的表达式为

.

(1)

其中：yi代表样本i的标签，正类为1，负类为0；pi代表样本i预测为正的概率.Dice系数CDice是一种对比图像相似度的度量函数，通常情况下用于计算两个图像数据的相似程度，它的取值范围为[0,1]，计算公式为

(2)

其中：N为图像数目；qi为模型预测中的某一像素；gi为相同位置上金标准中的某一像素.CDice越大表示其预测结果与真实结果的重合度越大，预测效果就越好.作为损失函数，其值越小，代表模型训练效果越好.Dice损失的表达式为

(3)

其中，FDice_Loss为Dice损失.本文将交叉熵与Dice损失的值相加后作为损失函数来评估模型训练的情况，其权重系数为1[21].

3 实验与分析

3.1 数据及数据预处理

本文采用的CSF图像信息数据集来自北京协和医院，采用HE染色技术获取，经过筛查得到29例可用数据作训练.这些训练数据集的标签是由专业人员手动标注而成，具有专业性.其标签的位深度为24 b，每个通道具有相同的像素数，且细胞像素设置为255，背景像素设置为0，所以标签的视觉效果是一个二值图像.

图像经过旋转、镜像、裁剪等方法将数据扩充至232例训练图像.并且将数据大小统一为512×512，以便训练.将扩充后的数据集经过一次拉普拉斯高通滤波处理，突出细胞边缘与背景的灰度差值，增强细节，其拉普拉斯核取为

3.2 环境配置

本文所使用的设备为i7-8700处理器，16 GB内存，64位Windows 10操作系统，NVIDIA GeForce GTX 1080 GPU加速处理图像.并且使用了TensorFlow-gpu 1.13.1与keras 2.1.5的开源深度学习框架，Python编程语言，Pycharm集成开发环境，还使用了Numpy计算库以及PIL中的图像处理方法等.

3.3 实验设置

将扩充后的数据经过预处理后传入网络中作训练.网络参数的优化器使用Adam optimizer[22]，此优化方法对高维度的数据进行快速拟合，利用梯度的一阶矩阵估计和二阶矩阵估计动态调整每个参数的学习率，经过偏置校正后，每次迭代学习率都有个确定范围，使得参数比较平稳.

为了均衡训练速度与训练效果，每次将2张图像传入网络中进行训练.加快网络训练速度，改善资源利用率，先将网络进行冻结，在冻结网络的情况下训练迭代50次，将更多的资源放在训练后面部分的网络参数.经过反复测试验证，冻结时的初始学习率设置为1×10-4效果最佳.迭代训练50次后，再将模型进行解冻，经过反复测试，这时的学习率设置为1×10-5效果最佳.模型解冻后继续迭代训练50次.当训练时，每隔3次当迭代模型性能不再提升时，则会采用每次将学习率降低原来学习率的0.5倍的方法来优化模型性能.

本文将实验数据按照8∶2的比例划分为训练集与验证集，将29例数据作为测试数据，来预测模型的实际分割效果.

3.4 评价指标与结果

本文采用了语义分割中3个最常用的评价指标进行分析分割结果，它们分别为相似系数EDice，平均交并比U，以及类别平均像素准确率A.本文以二分类为例，计算公式为

(4)

(5)

(6)

其中：TP，FP和FN分别为模型预测的真正例、假正例和假反例；P1和P2分别为类别1和类别2的像素准确率，定义为

(7)

(8)

经过所有网络的训练，选择两张具有代表性的CSF细胞图像进行预测分割展示.图8为细胞大小一般且排列相对稀疏的数据图，图9为细胞较小且排列密集的数据图.

图8 细胞稀疏图像

图9 细胞密集图像

如图所示，使用传统Otsu阈值分割时，属于背景的部分像素会被判定为细胞，对噪声极其敏感，对灰度差异不明显及不同目标灰度值有重叠的分割不明显.而PSPnet网络对医学图像的细节还原性较差，无法很好地识别医学图像的小目标物体.Segnet与DeeplabV3+模型分割尺寸相对差不多大小的细胞时，会存在细胞黏连的问题，并且分割尺寸相对较小且数量较多的细胞时，存在细胞分辨不清晰、分割模糊的问题.U-Net分割模型基本上实现了对所有细胞的有效分割，但是依然存在不足之处，该网络模型对细胞边界与背景之间的区分度略有不足.而使用本文的分割模型做预测分割时，上述所有问题均得到了改善.

使用上述分割模型分别对CSF细胞图像做出预测，EDice，U和A的具体结果如表1所示.

表1 CSF分割结果

由表1可看出，使用本文方法对CSF细胞图像做细胞分割要优于其他分割方法.

为了验证本文方法的普适性，选用2018 Data Science Bowl公开数据集作为参照实验数据.该数据包含大量分割的核图像，图像是在各种条件下获得的，细胞类型、放大倍数和成像方式各不相同.使用此数据得到预测分割的各项指标具体结果如表2所示.

表2 2018 Data Science Bowl数据集分割结果

由表2可看出，在2018 Data Science Bowl数据集分割中，本文方法要优于其他分割方法.

4 结 论

本文以CSF细胞为分割研究对象，采用基于注意力机制的U-Net网络，经过数据扩充与预处理，引入迁移学习，提高资源利用率，并且与Otsu，PSPnet，Segnet，DeeplabV3+以及U-Net做了对比实验，结果表明，本文方法要优于其他分割方法.并且通过公开数据集2018 Data Science Bowl做了进一步对比实验，实验结果进一步验证了本文分割方法在各项指标上均优于其他分割方法.