基于CBA-YOLO模型的煤矸石检测
2022-07-07桂方俊李尧
桂方俊, 李尧
(中国矿业大学(北京) 机电与信息工程学院,北京 100083)
0 引言
煤矸石分选是煤炭洗选加工的重要环节,主要有人工拣矸法、湿选法和干选法3种方式。人工拣矸法效率低,湿选法存在水资源浪费等问题,干选法是目前煤矸石分选领域的研究热点[1],符合煤矿绿色发展的要求。煤矸石识别技术是煤矸石干选的核心技术,其原理是利用提取的样本特征判断当前被测物是煤还是矸石。煤矸石识别方法包括传统方法和基于图像处理的方法:① 传统方法有放射性探测法、密度识别法等。放射性探测法根据煤和矸石对射线吸收程度不同进行识别,需要较高的执行速度,一定程度上受煤矸石含水量的影响,并且存在射线辐射问题。密度识别法通过光电技术获取被测物质量和厚度,计算被测物密度,以此识别当前被测物为煤或矸石,该方法流程复杂且对设备要求高。② 基于图像处理的方法成本较低、便于推广,近年来应用广泛,机器学习和深度学习是其中的两大分支。机器学习方法通过图像特征提取和分类算法识别目标,对应用场景的要求高,鲁棒性差。深度学习方法通过提取图像高维特征并进行推理,鲁棒性好,准确率高,因而在煤矸石检测领域被逐步推广。
将深度学习应用到煤矸石检测中的研究已取得一定成果。文献[2]基于LeNet-5进行改进,实现了煤矸石图像检测和分类。文献[3]通过AlexNet卷积神经网络模型提高了煤矸石检测的准确率。文献[4]通过CornerNet-Squeeze深度学习模型提取特征,利用图像灰度直方图的三阶矩特征参数对煤矸石进行分类,准确率显著提升。上述方法虽然实现了煤矸石检测,但在性能上与超快速的区域卷积神经网络(Faster Region-Convolutional Neural Network,Faster R-CNN)[5]、单阶段多框检测器(Single Shot MultiBox Detector,SSD)[6]、YOLO(You Only Look Once)系列[7]算法相比还有一定差距。以Faster R-CNN系列为代表的两阶段算法对区域建议网络和主体检测网络分别进行训练,准确率高,但实时性表现不佳。以SSD,YOLO系列为代表的单阶段算法将种类识别和目标定位视为一个回归问题来解决,建立端到端训练,具有较高精度,同时保证了实时性。文献[8]将最新的YOLOv5算法改进后应用到煤矸石检测中,通过提升模型的特征提取能力达到优化效果,但并未针对煤矸石检测任务中的具体特性进行分析优化。
针对煤矸石检测中存在的特征差异不大、目标密集等问题,本文在YOLOv5基础上,通过卷积块注意力模块(Convolutional Block Attention Module,CBAM)、双向特征金字塔网络(Bi-directional Feature Pyramid Network,BiFPN)结构、Alpha-IoU函数对模型的各部分进行改进,提出了一种基于CBA-YOLO模型的煤矸石检测方法,通过实验验证了该方法的有效性。
1 YOLOv5基础模型选择
YOLO是一种基于深度神经网络的目标检测与定位模型,其最大特点是运行速度快,可在实时系统中使用[9]。YOLOv5是目前YOLO系列的最新版本,继承了YOLOv4[10]算法的优势,同时在主干网络中引入跨阶段局部(Cross Stage Partial,CSP)结构[11],进一步提升了图像推理速度,模型结构更加小巧。YOLOv5共有4个版本,依据网络大小排序分别为YOLOv5s,YOLOv5m,YOLOv5l,YOLOv5x,4种模型的性能如图1所示。
图1 YOLOv5模型的性能Fig. 1 The performance of YOLOv5 models
图1 中,曲线越靠近左上角区域,表示模型的精度和速度越高。由图1可知,YOLOv5s模型的速度快但精度低,YOLOv5x模型的精度高但速度最慢,YOLOv5m和YOLOv5l模型速度和精度很均衡,但YOLOv5l模型更大,综合考虑算力、精度和速度,本文选用速度较快、精度较高的YOLOv5m为基础模型。
2 CBA-YOLO模型
YOLOv5模型主要由输入端、Backbone、Neck、Prediction组成。CBA-YOLO模型在YOLOv5的基础上进行了改进:在Backbone中加入CBAM,以聚焦特征差异,降低数据维度,提升检测性能;将Neck部分改为BiFPN结构,以提高检测速度;在Prediction部分,采用Alpha-IoU作为损失函数,以提高检测精度。
2.1 输入端
在实际应用中,数据集通常包含大量小目标,小目标分布不均匀的问题很突出。YOLOv5采用Mosaic数据增强方式进行数据处理,针对小目标检测问题,将4幅图像以随机缩放、随机剪切、随机安排的方式拼接在1张图像里,以达到丰富数据集的效果。这种随机处理的方式增加了小目标数量,增强了模型的鲁棒性,同时减少了对GPU的占用,1个GPU就可获得较好效果。
2.2 Backbone改进
YOLOv5的Backbone由多个卷积模块堆叠而成,缺点是无法聚焦特征差异,而煤矸石检测任务中,煤和矸石目标特征差异不明显。因此,在Backbone中加入CBAM进行改进,以聚焦特征差异。
注意力机制是深度学习中的一种数据处理方法,主要作用是忽略无关信息而关注重点信息,提升信息处理的效率和准确性。注意力机制通常分为空间注意力机制[12]、通道注意力机制[13]、空间和通道混合注意力机制[14]。CBAM属于空间和通道混合注意力机制,其结构如图2所示。
图2 CBAM结构Fig. 2 Structure of CBAM
通道注意力模块同时使用最大池化和平均池化,将多层感知器(Muti-Layer Perception,MLP)层分别应用于2个通道,经过sigmoid激活函数合并通道权重,最后逐通道加权得到attention结果。空间注意力模块在通道重定义特征基础上进行最大池化和平均池化,得到2个单通道的特征图,再通过一个7×7的卷积形成特征图,最后逐通道加权获得attention结果。这2种模块分别学习了通道的重要性和空间的重要性,CBAM将其串联起来,可产生更有分辨性的特征表示,同时可降低数据维度,加速网络推理过程。
改进Backbone结构如图3 所示。原始大小为640×640×3的图像输入Focus结构,采用切片操作转换成大小为320×320×12的特征图,再经过一次32个卷积核的卷积操作,转换成320×320×32的特征图后输入CBAM模块。通过通道注意力模块获得通道注意力特征,确定每层特征图上需要关注的重点内容;再通过空间注意力模块获得空间注意力特征,确定需要关注的重点位置;最后将串联结果输入后续的跨阶段局部网络层(Cross Stage Partial Network,CSPNet) 和空间金字塔池化(Spatial Pyramid Pooling,SPP) 结构中处理。
图3 改进Backbone结构Fig. 3 Structure of improved Backbone
2.3 Neck改进
YOLOv5的Neck采用特征金字塔网络(Feature Pyramid Network,FPN) + 路 径 聚 合 网 络(Path Aggregation Network,PAN)结构,如图4(a)所示。该结构通过PAN增加一个自底向上的路径聚合网络,解决了传统的自顶向下的FPN受单向信息流限制的问题。但是由于煤矸石检测任务目标密集,单张图像中目标众多,计算量大,导致YOLOv5模型应用于煤矸石检测任务时实时性不佳。为了保证煤矸石检测任务的实时性,在Neck中引入BiFPN结构[15],如图4(b)所示。
图4 特征网络结构Fig. 4 Structure of features network
改进方法:① 去除只有1条输入边的节点,得到1个简化的双向网络,若1个节点只有1个输入,没有特征融合,则其对以特征融合为目标的特征网络的贡献较小。② 当原始输入节点与输出节点处于同一水平时,增加1条从原始输入节点到输出节点的额外边,以在不增加成本的情况下融合更多特征;③ 将每个双向(自顶向下和自底向上)路径视为1个特征网络层,并将多个特征网络层进行叠加,以实现更高级的特征融合。
2.4 Prediction改进
YOLOv5的Prediction由损失函数和非极大值抑制组成。损失函数用于衡量模型生成的预测框与真实框之间的重叠程度,交并比(Intersection over Union,IoU)[16]是一种常用的损失函数。非极大值抑制用于目标检测的后处理过程中,针对多目标框的筛选,抑制非极大值元素,搜索局部极大值,去除冗余的检测框,得到最终预测结果。YOLOv5的损失函数为GIoU[17],但是当预测框和真实框出现包含现象时,无法达到优化效果。文献[18]采用CIoU作为损失函数,通过计算检测框之间的欧氏距离而非其IoU,解决两框包含时出现的问题,同时增加了检测框尺度、长和宽的损失,使预测框与真实框更相符,但其结果相对固定,不能自适应地对高IoU目标和低IoU目标的损失和梯度进行加权。因此,本文引入权重系数α,用Alpha-IoU[19]替代CIoU作为预测目标的边界框损失函数,其计算公式为
式中:I′为Alpha-IoU的值;I为IoU的值;ρ2(b,bgt)为预测框中心点b与真实框中心点bgt的欧氏距离;d为可以包含预测框和真实框的最小闭合区域的对角线长度;β为trade-off 的参数;γ为衡量框的长宽比一致性的参数;l为损失值。
根据目标的IoU值自适应地调整权重系数α,有助于检测器更快地学习高IoU目标,从而提高目标检测精度。
3 基于CBA-YOLO模型的煤矸石检测
基于CBA-YOLO模型的煤矸石检测流程如图5所示。输入图像或视频,对视频进行逐帧读取并转换为图片格式;通过自适应算法将图像尺寸统一调整为640×640,输入CBA-YOLO模型;通过Backbone、Neck、Prediction进行计算处理,得到煤矸石的相关信息,以数组的方式输出检测结果,数组结构为[检测框左上坐标,检测框右下坐标,类别,置信度]。
图5 基于CBA-YOLO模型的煤矸石检测流程Fig. 5 Flow of coal gangue detection based on CBA-YOLO model
4 实验分析
4.1 数据采集
为了验证本文方法的有效性,在河北某煤矿选煤厂的选矸装置上架设Blaser工业相机进行数据采集,并通过大功率LED面光源和光源控制器提供光照均匀、稳定的照明环境,如图6所示。
图6 图像采集Fig. 6 Image acquisition
设置每隔2 ms采集1张图像,共采集1 500 张图像。为了达到更好的网络训练效果,选取1 245张质量较好的图像作为数据集,用LabelImg标注软件对数据集进行标注,按照8∶1∶1的比例划分为训练集、评估集和测试集3个部分,其中训练集995张,评估集125张,测试集125张。数据集中包含单块煤矸石图像和多块煤矸石图像。
4.2 CBA-YOLO模型训练
实验软硬件配置如下:计算机操作系统为Ubuntu16.04,CPU运 算 核 心 为Intel®CoreTMi5-8400CPU@2.80 GHz,GPU为GeForce GTX 2080Ti,NVIDIA驱动为CUDA10.2。算法基于Pytorch框架运行。实验训练参数设置如下:选用ADAM作为优化器;初始学习率为0.01,学习率周期为0.2,学习率动量为0.937;权重衰减系数为0.000 5,训练的Batch_size为16,epoch为600。
对CBA-YOLO模型进行多次训练,选取最佳模型用于评估。训练完成后的损失函数曲线如图7所示。损失值越小,预测结果越准确。由图7可知,训练轮次超过250后,损失值下降趋于平缓,损失函数曲线已经收敛,说明模型训练效果良好。
图7 训练损失Fig. 7 Training loss
4.3 实验结果分析
利用获得的测试集进行消融实验,实验1-3分别在模型中引入CBAM,BiFPN结构和Alpha-IoU,实验4-6采用2种改进方法。为了简洁表达,取CBAM,BiFPN,Alpha-IoU的首字母,将添加单个CBAM模块的模型称为YOLO-C,添加CBAM模块和BiFPN组合结构的模型称为YOLO-CB,其他模型依此类推。
通过设置不同阈值,得到不同模型准确率与召回率的关系曲线,即PR曲线,如图8所示。PR曲线与坐标轴围成的区域面积为平均精度均值(mean Average Precision,mAP),以mAP作为煤矸石检测的精度指标,以帧率作为煤矸石检测的速度指标,消融实验结果见表1。
图8 消融实验PR曲线Fig. 8 PR curves of ablation experiment
表1 消融实验结果Table 1 Results of ablation experiment
分析图8和表1可知,以YOLOv5为基准模型,采用单个改进方法的模型性能如下:YOLO-C的精度和速度都有小幅提升;YOLO-B的精度小幅提升,速度提升显著;YOLO-A的精度提升最多,但速度有所下降。上述结果表明:在Backbone中加入注意力机制可聚焦特征差异,降低了数据维度,带来了精度和速度的提升;Neck部分引入BiFPN结构后,提升了模型计算效率,但精度提升不明显;Prediction部分通过改进损失函数,聚焦高IoU目标,提升了检测精度,但未能兼顾速度。
以YOLOv5为基准模型,采用组合改进方法的模型性能如下:YOLO-CA和YOLO-BA的精度提升程度相当,但后者的速度更快,说明Neck部分采用BiFPN结构可提升网络推理速度;YOLO-CB的精度最低,但速度最快;CBA-YOLO模型的精度提升最多,比基准模型提升了3.4%,检测速度提升了10%,说明改进损失函数虽然增加了网络计算量,但由于Neck部分采用BiFPN结构后提升了模型计算效率,保证了模型在不损失速度性能的情况下稳定提升精度,同时,通过在Backbone中加入注意力机制进一步提高了检测性能。
基于YOLOv5和CBA-YOLO模型的煤矸石检测结果对比如图9所示。可看出,CBA-YOLO模型鲁棒性更强,有效避免了漏检、误检和重叠现象。
图9 煤矸石检测结果对比Fig. 9 Comparison of coal gangue detection results
5 结论
(1) 在Backbone中加入注意力机制可聚焦特征差异,降低了数据维度,带来了精度和速度的提升;Neck部分引入BiFPN结构后,提升了模型计算效率,但精度提升不明显;Prediction部分通过改进损失函数,聚焦高IoU目标,提升了检测精度,但未能兼顾速度。
(2) 与YOLOv5模型相比,CBA-YOLO模型的精度提升了3.4%,检测速度提升了10%,同时保证了检测实时性和精度,可为选煤厂实际应用提供参考。