基于自注意力机制的矿井次光照图像语义分割研究

杨克虎,龙启航,汪嘉文,彭宝山,金 波,杨学孟

(1.中国矿业大学(北京) 机电与信息工程学院,北京 100083; 2.国网能源新疆准东煤电有限公司,新疆 准东 831599;3.应急管理部煤矿智能化与机器人创新应用重点实验室,北京 100089;4.中国矿业大学(北京) 内蒙古研究院,内蒙古 鄂尔多斯 017010)

煤矿智能化建设,是整个煤炭行业高质量发展的必由之路[1]。其中,对煤炭开采场景的准确理解是煤炭智能开采的基础。语义分割作为图像场景理解的基础任务,能够对煤炭开采图像场景实现像素级别的分类,准确分割出巷道壁、无人矿车、障碍物、井下设备、人员等目标物,有助于工作人员掌握井下工况,避免设备相撞等安全事故的发生。

常用的标准语义分割数据集有PASCAL VOC[2]、COCO[3]和Cityscapes[4]。这些数据集包含了室内外大部分常见场景,精细标注了上千张图片,涉及上万个标注实体,正是这些标准数据集的发展,使得神经网络有充足的数据可使用,才能实现对室内外大部分常见场景的精准分割。

在常见场景中,应用于语义分割的卷积神经网络已获得一定的发展。FCN网络由LONG等提出,首次将应用于图像分类神经网络的全连接层换成卷积层,使神经网络能对输入的任意尺寸图像进行处理,完成语义分割[5];U-Net网络[6]采用编码器—解码器的网络结构,对图像局部区域的提取特征效果好;DeepLab系列网络[7]实现了多尺度特征提取;为了追求更快的算法速度,学者们提出了具有双边分割网络结构的BiseNet[8],可应用于自动驾驶等领域。

近年来,随着自注意力机制的兴起,越来越多的网络开始结合自注意力机制进行语义分割,并获得了良好效果。将Transformer网络[9]与传统语义分割网络结合,XIE等[10]提出了新型网络结构Segformer,在获得更高精度的同时,大大提升了网络的实时性。神经网络的不断发展,使语义分割技术的精确性、实时性日渐提高,开始应用于自动驾驶、医学图像等多个领域。

针对次光照矿井图像的分割,不少专家学者提出了自己的方法。单鹏飞等[11]提出了基于改进Faster R-CNN的方法来进行煤矸识别,最终达到了86.53%的平均像素精度;薛雯予等[12]提出了对矿井巷道图像使用中值滤波和改进Top-hat滤波法,去除图像噪声,利用迭代算法,求取图像分割的阈值,分割出井下图像目标区域,获取目标点坐标来增强图像;韩斌等[13]采用基于改进CV模型火灾图像分割方法,分割识别出矿井早期火源发生点位;高峰等[14]提出了性能优于传统语义分割网络的一种煤岩图像分割识别网络模型来分割煤岩图像;李晓宇等[15]使用基于超像素特征与SVM分类的人员安全帽分割方法,对矿井图像中人员所佩戴的安全帽进行分割。

上述研究虽然在矿井图像分割精度上取得了丰硕成果,但是分割种类均为二分类,分割实例种类过少,数量远不及普通场景语义分割。为此,杨潇等[16]提出基于域适应的煤矿环境监控图像语义分割,将标准数据集视为源域数据集,将采集的煤矿巷道图片作为目标域数据集,利用双对齐网络结构,进行域适应语义分割,最终在包含有6种实例的煤矿数据集上得到38.26%的平均交并比。但此方法分割精度较低,难以直接应用于实际煤矿安全生产中。

针对采集到的矿井暗光图像,笔者使用基于深度学习的图像增强算法进行处理,增强图像边缘,再利用LabelMe标注出连同背景在内的11种类别,专门构建了标准矿井图像语义分割数据集,并提出了基于自注意力机制的轻量级编码—解码结构卷积神经网络,进行矿井图像的语义分割。对比实验证明,该种网络能够做到准确分割,与主流语义分割网络相比,更能提取次光照矿井图像特征,鲁棒性较好。

1 矿井图像语义分割数据集构建

采集了河北唐山、内蒙古鄂尔多斯、北京密云三地的矿井巷道图像,挑选出具有代表性的16种主要矿井图像场景。拍摄矿井图片,采用微光图像增强算法增强图像边缘信息,并使用LableMe(标注软件)对图像进行人工标注,同时为了避免网络过拟合,利用单应变换矩阵对图像进行单应变换,扩充数据集。矿井巷道图像数据集构建过程如图1所示。

图1 矿井巷道图像数据集构建过程

1.1 图像增强

矿井图片如图2所示。实验需对井下图像中的巷道壁、轨道、铁轨、管道、灯、障碍物、矿车、人、货车、指示牌共10类场景进行标注。

图2 矿井巷道图像数据集构建

由图2可以看出,因矿井大部分场景图像拍摄于照明条件较差的环境下,巷道图像整体色彩偏灰暗、像素值偏低、图像边缘信息不明显,图中的目标物体用肉眼难以分辨。这会给后续的人工标注带来极大困难。为此,使用微光图像增强算法[17]对次光照下拍摄的图像进行增强, 增强前后对比如图3所示。

(a) 原图

(b) 微光图像增强后的图

微光图像增强算法是一种基于深度学习的图像增强方法:在轻量级网络DCE-Net结构中输入图像,输出用于调整图像像素的高阶曲线,并利用高阶曲线对原图像的像素进行一一映射,对像素值进行动态调整。此算法应用于矿井暗光图像上能取得良好的图像增强效果。由图3(b)可以看出,巷道较暗场景经过微光图像细节增强后,地面上的铁轨、沙砾清晰可见。

1.2 图像标签掩码生成

选取的矿井巷道图像数据集包含1 270×720和640×480两种分辨率。综合分析各场景图像,选取了连背景在内的11类场景,使用LabelMe对图像进行精细标注。对940张经去雾算法和微光图像增强算法处理后的图像,以及1 699张原始清晰图像,共计2 639张精细标注的图像进行标注,一共标注了26 521个对象。

场景的具体实例如图4所示,标签的具体名称和序号见表1。

(a) 原图

(b) 标注掩码后的图

表1 标签信息

每个对象的数量统计见图5。其中,巷道壁虽然占据了图像的大部分区域,但整体数量偏少;灯虽然占据图像区域小,但数量偏多。因此矿井巷道数据集对卷积神经网络的局部区域和全局区域分割均提出较高要求。

图5 数据集每个对象数量统计

训练相同的轮数,数量偏少的种类可能会出现训练不完全的现象。为了增加偏少种类的数量,同时也为了避免过拟合问题的产生,对图像和掩码标签同时采取单应变换处理以扩充数据集:将图像按随机生成的角度、尺度,分别进行旋转、缩放(见图6),将原始数据集扩充至9 996张图像。

图6 数据集扩充过程

2 基于自注意力机制的轻量级编码—解码网络

2.1 整体网络结构

从实际生产角度出发,模型需要在分割精度和推理速度上均达到较好效果。以DeepLab V3+[18]网络的编码—解码结构为基础,构建基于自注意力机制的轻量级编码—解码网络(以下简称本文算法):在编码结构中结合自注意力机制和ShuffleNet v2轻量级网络[19],利用池化金字塔结构进一步提取语义特征信息;在解码结构中,使用2次线性插值上采样恢复原特征图尺寸,最终输出语义分割结果。网络整体结构如图7所示。

图7 基于自注意力机制的编码—解码网络模型

编码结构用来提取图像的深层和浅层语义信息,其工作原理如下:

1)利用深度可分离卷积和通道信息融合层构建网络骨干,提取输入的任意尺寸RGB图像的语义特征信息,输出通道数为464的深层语义特征信息和通道数为24的浅层语义特征信息。

2)在主干提取网络之后,深层语义特征经由自注意力将模块特征激活后输入到并行分支ASPP结构中,进行加强特征提取。ASPP是由5个分支并行组成的加强提取特征结构:使用了1个1×1的卷积、3个3×3的膨胀卷积和1个全局池化层来对深层语义特征进行进一步的加强特征提取。将经过5个分支得到的输出层在通道方向上进行拼接,并由1×1卷积调整通道数。

3)将处理后的深层语义特征信息、浅层语义特征信息输入到解码结构中。解码结构中,深层语义特征信息在进行采样后恢复至和浅层语义特征信息同样尺寸大小,在通道方向上和浅层语义特征信息拼接后再进行整体采样,输出分割结果。

2.2 语义特征信息提取

DeepLab V3+使用的原始特征提取网络是Xception[20]结构。Xception由多个并行分支组成。虽然其在一定程度上很好地融合了特征层的空间信息,但因拥有更深的网络结构,对算力要求更高。为了使网络结构更加轻量化,同时提高模型分割准确率,使用通道信息融合结构和深度可分离卷积构建的网络主干,输出提取的深层语义特征信息和浅层语义特征信息[19]。用于语义特征信息提取的基础模块block1、block2如图8所示。block1中,输入特征层首先经过Channel Split,然后被均分成两部分,一部分不经过任何操作,另一部分经过由1×1卷积、步距为1的3×3DW卷积、1×1卷积和BN层、ReLU激活函数构成的分支,2个分支的输出在深度方向拼接在一起,再经过Channel Shuffle层融合得到最终的输出特征。

Channel Shuffle融合不同通道之间的信息,其工作原理如图8(b)所示。输入三通道Channel1、Channel2、Channel3,将每个通道的特征均分成3份,分别将其中的第1份取出来放在一起,第2份和第3份也进行同样的操作,最终得到Channel Shuffle层。

(a) block结构

(b) Channel Shuffle

将block1(步距为1)、block2(步距为2)、卷积层和池化层进行堆叠,构成主干语义特征提取网络,并提取深层语义特征信息和浅层语义特征信息,如表2所示。

表2 主干网络结构

1)输入,Conv1。矿井巷道图像先经过卷积核为3×3、步距为2的卷积2倍下采样,并将通道数设为24,此时输出浅层特征信息。

2)MaxPool。经过步距为2的最大池化层,对输入进行2倍下采样,但输出通道数不变。

3)Stage2。block2重复1次,block1重复3次,并将通道数调整为116。

4)Stage3。block2重复1次,block1重复7次,并将通道数调整为232。

5)Stage4。block2重复1次,block1重复3次,并将通道数调整为464,输出深层特征信息。

2.3 自注意力机制计算

在编码—解码的语义分割网络中,编码过程往往只利用到了局部信息,而忽略了全局信息。为了实现对全局信息的考察,引入了自注意力机制计算模块。在视觉任务中,通常使用Transformer模块计算自注意力机制。Transformer模块由Layer Norm层、多头自注意力机制计算层(Multi-Head Attention)和MLP模块构成。MLP模块又由简单的线性层、GELU激活函数和Dropout层构成,其流程如图9所示。

多头自注意力机制计算层是Transformer模块的核心。将展平后的二维向量通过Wq、Wk、Wv3个变换矩阵得到相应的特征向量Q、K、V,计算得到每一个像素的权重系数,与输入特征矩阵相乘,完成自注意力机制的计算,计算公式如下:

(1)

但由于每一个像素都会通过计算得到一个权重值,这就对计算机的算力提出了较高要求。笔者将MobileVit[21]中的Unfold、fold层与2.2节轻量化语义特征信息提取主干相结合,在提高分割精度的同时,也满足分割速度的落地需求。将通道数为464的深层语义特征信息输入自注意力计算模块中进行特征激活。在此过程中会进行Unfold和fold层,这2个层能有效减小计算量。将特征图划分成多个patch,再进行自注意力机制的计算,如果patch大小划分为2×2,只对图中颜色相同的部分进行自注意力机制计算,例如4个红色小块堆叠成特征向量,进行自注意力机制的计算(见图10),再通过fold过程还原成原来结构(见图11)。

图9 Transformer结构

图10 自注意力计算

图11 Unfold过程示意图

此时计算量P1:

(2)

式中W、H、C分别为特征图的宽、高和通道数量。

在一张图中,每个像素点周围的像素值具有相似性,因此没有必要将每个像素都与其他每个像素进行计算,否则会大大增加计算量。计算量P2:

P2=WHC

(3)

综上所述,Unfold层和fold层能大大减少计算量并提高网络模型精度。

2.4 网络解码结构

解码结构中,深层语义特征信息和浅层语义特征信息经过卷积调整通道数目,在深度方向上进行拼接后,输出最终的分割结果。

本网络采用的交叉熵损失函数计算公式如下:

(4)

式中:y为交叉熵;yc为前景或背景标识,取1或0;Pc为经过softmax之后的概率值。

3 对比实验

3.1 实验过程

采用的评价指标为每一类别交并比IoU(%)、每一类别像素精度PA(%)、平均交并比mIoU(%)和平均像素精度mPA(%),综合记录在P100显卡上推理一张图片的速度。IoU计算公式如下:

(5)

式中:ti为类别i的总像素个数;nii表示类别为i、预测值也为i的像素个数;nji表示类别为j、预测值为i的像素个数。

全部像素种类的平均交并比计算公式如下:

(6)

式中ncls为总的像素个数。

每一类预测正确的像素个数占此类像素总个数的比例即为像素精度,其计算公式如下:

(7)

每一类别的像素精度求和取平均值即为全局像素精度,其计算公式如下:

(8)

将矿井图像数据集按照8∶2的比例划分为训练集和验证集。实验使用2张NVIDIA Tesla P100显卡进行训练。在PASCAL VOC2012上进行预训练,得到预训练权重,将预训练权重加载至网络中,再利用矿井图像数据集进行训练。在矿井巷道图像数据集上训练时,先冻结网络主干,训练解码部分,再对网络进行微调。 将batch_size设置为4,采用SGD优化器,优化器内部的momentum参数设置为0.937,为了防止网络过拟合,权值衰减系数设为5×10-4,优化器的初始学习率设置为1×10-2,由于巷道壁占据图像的绝大区域,将巷道壁类别的初始权重系数设置为1.00,其余类别初始权重系数设置为0.50,使用交叉熵损失作为损失函数。

选取同样是轻量级的网络BiSeNet[21],由轻量级网络MobileNet[22]改进的DeepLab V3+、PSPNet,以及由ResNet网络改进的PSPNet[23]网络进行实验对比。在所有网络模型中输入1张512×512的彩色图像,用网络模型参数数量和浮点运算数量FLOPs初步对比网络模型参数和计算量,判断网络结构是否轻量化。

3.2 实验结果

计算量与参数,精度与推理速度的对比结果见表3～4。

表3 各网络计算量及参数

由表3可知:理论计算量FLOPs方面,本文算法为48.80 G,低于PSPNet(ResNet)和DeepLab V3+(MobileNet),高于PSPNet(MobileNet)和BiSeNet;模型参数方面,本文算法为11.90 M,低于PSPNet(ResNet)和BiSeNet,高于PSPNet(MobileNet)和DeepLab V3+(MobileNet)。

由表4可知:精度方面,本文算法mIoU表现最优,为76.50%,比PSPNet(ResNet)高2.44%,比PSPNet(MobileNet)高28.78%,比DeepLab V3+(MobileNet)高1.65%,比BiSeNet高5.77%;本文算法mPA比PSPNet(ResNet)高5.28%,比PSPNet(MobileNet)高28.68%,比DeepLab V3+(MobileNet)高3.17%,比BiSeNet高5.28%;在推理速度上,以MobileNet为backbone的PSPNet和BiSeNet分别能达到0.022、0.030 s,本文算法为0.032 s,略高于传统轻量级网络BiSeNet的推理速度。

表4 各网络训练精度及推理速度

具体在某一类标签上取得的像素精度和交并比,如表5～6所示。

表5 各网络在每一类别上的像素精度

表6 各网络在每一类别上的交并比

由表5～6可见,本文算法在11种类别的像素精度和交并比方面均取得优异的表现,在背景、巷道壁、轨道、铁轨、管道、灯、障碍物、矿车、人、货车和指示牌等类别中,像素精度分别能达到94.76%、96.25%、88.68%、73.75%、81.49%、73.34%、81.49%、97.80%、95.41%、97.19%、85.12%,交并比分别能达到90.07%、93.45%、78.73%、58.67%、66.46%、65.16%、60.32%、93.00%、87.99%、89.54%、58.13%。

综上所述,本文算法在平均交并比、平均像素精度,以及每一类别交并比、像素精度上均能取得优异的表现,在推理速度上也略高于传统轻量级网络BiSeNet。无论在占据图像区域面积较小但总体数量较多的类别中,还是在占据图像区域面积较大但总体数量较少的类别中,本文算法的像素精度和交并比均优于其他网络,综合性能优于其他网络。

3.3 图像分割效果

本文算法与DeepLab V3+(MobileNet)方法的mIoU较为接近。为了更进一步证明其优越性,使用包含对象类别较多、较为复杂的场景图像对这2种方法预测效果进行测试。结果表明,本文算法不论是对全局信息还是对局部信息的提取上均要比以MobileNet为主干提取特征信息的DeepLab V3+网络表现良好。复杂图像推理效果对比如图12所示。

(a) 原图 (b) DeepLab V3+(MobileNet)(c) 本文算法

为了更加明显地展示本文算法的预测效果,分别对原图1～4进行预测,各种算法预测效果如图13所示。

图13 各种算法预测效果对比

由图13对比可以看出,本文算法在全局信息和局部信息上均能得到很好的推理效果,能准确分割图13(a)中地上的轨道、铁轨和障碍物的细节轮廓;能精确分割细微的远景如图13(b)中主要包含近处的工作人员,但整体图像偏黑,远处墙上的水管和轨道肉眼难以区分;能分割出图13(c)中墙上管道的整体轮廓;图13(d)为旋转过后的图像,其中的巷道壁的语义信息难以提取,尤其是巷道壁边缘信息,PSPNet不能准确提取出图13(d)巷道壁,但本文算法能准确分割提取,保障了生产安全。因此,本文算法在网络复杂度、推理速度、分割精度上都有良好的表现。

4 结束语

1)采集多地矿井巷道图像,针对次光照图像进行图像细节增强,并标注出人员、设备等10种类别,构建了标准的矿井语义分割图像数据集。

2)提出基于自注意力机制的编码—解码结构网络,以DeepLab V3+网络为基础,在编码器部分加入轻量级自注意力机制模块进行特征信息提取,在解码器部分对深、浅层语义特征信息进行拼接,恢复原始图像尺寸,输出语义分割结果。

3)针对3通道512×512像素大小的图像,基于自注意力机制的编码—解码结构网络的理论计算量FLOPs为48.80 G 、模型参数量为11.90 M;在P100显卡上,推理速度能达到0.032 s/张;在分割精度上mPA为87.75%、mIoU为76.50%,优于PSPNet和以MobileNet为主干的DeepLab V3+、BiSeNet网络;在具体类别像素精度上,矿车的分割精度能达到97.80%,货车的分割精度能达到97.19%;而在每一类别的交并比表现上,巷道壁交并比能达到93.45%。分割准确率高,鲁棒性强。