基于改进YOLOv5的矿石目标检测算法

Ore Target Detection Algorithm Based on Improved YOLOv5

摘要: 针对工业矿山生产线上工作量大，在进行目标定位检测时需兼具较高实时性，提出一种基于改进YOLOv5的目标检测算法。该算法在主干特征提取网络引入更轻量化的MobileNetv3，同时将PANet网络3 × 3的卷积块改进，提高了检测精度和效率；使用DIoU-NMS替代原模型中NMS的方式，提高矿石重叠时检测精度。改进后的算法的参数量减少了24.95%，mAP值达到了95.7%，帧数达到了166.7 FPS，检测单张图片的时间缩短了4 ms，表明该算法在矿石传送带上具有更强实用性和更高的检测精度。

Abstract: Aiming at the heavy workload in the industrial mine production line and the high real-time performance in the target positioning detection, a target detection algorithm based on improved YOLOv5 is proposed. The algorithm replaces the backbone feature extraction network CSPDarknet with the lighter MobileNetv3, and improves the 3 × 3 convolutional block of PANet network to improve the detection accuracy and efficiency. Using DIoU-NMS to replace the weighted NMS in the original model further reduces the number of parameters, improves the convergence of the model, and improves the detection accuracy when the targets overlap. The parameter number of the improved algorithm is reduced by 24.95%, the mAP value is 95.7%, the frame number is 166.7 FPS, and the detection time of a single image is shortened by 4 ms, which shows that the algorithm has stronger practicability and higher detection accuracy on the ore conveyor belt.

1. 引言

随着信息技术和机器人技术的发展，工业生产领域对自动化程度要求越来越高。而对于大型的钨矿石原料，一般利用颚式破碎机将矿石进行破碎，得到较小粒度的矿石。在矿石破碎磨矿加工过程中，往往必不可少需要装配有胶带运输设备，主要用于运输中矿石。输送的矿石中夹杂有超过限制要求的较大块状矿石。而现有技术中，由于各种原因，目前较多采用手工手选的方式进行分选。如不能及时进行分选，大块矿石进入进料口造成堵塞，强行进行破碎可能会引起设备工作过载、机器剧烈震动、设备内部构造磨损甚至损坏。因此将现今热门的机器视觉技术结合工业机器人技术应用于矿石分选 [1] ，在准确识别定位需分拣出的大块矿石后用机械臂将其带离传送带。

目标检测技术已然成为当下对于目标定位抓取领域最热门的研究方向。目标检测方法大致分为两类：以YOLO [2] 系列为代表的单阶段和以RCNN [3] 为代表的两阶段。张艳等提出利用金字塔分割注意力模块，捕获不同尺度特征图的空间信息以丰富特征空间，提高目标检测的精度，利用基于线性变换的Ghost Bottle Neck模块有效减少模型参数量 [4] 。王建才等将CABM注意力机制和Transformer添加到YOLOv5模型中，大大提高了煤岩图像的检测精度 [5] 。冒国韬等利用在骨干网络瓶颈层的多尺度分割注意力单元建立多尺度特征间的远程依赖关系，强化关键特征的表达能力并抑制背景噪声干扰，在复杂背景下仍能提升图像目标检测性能 [6] 。

综上，针对矿山传送带分拣大块矿石智能化检测定位困难的问题，本文基于YOLOv5模型对大小不一的矿石进行识别。同时，为了提高复杂背景下矿石图像检测效果，本文采取基于改进YOLOv5算法对矿石图像进行识别，为后续实现工业机械手实时分拣提供基础。

2. YOLOv5神经网络

YOLOv5网络结构主要由三个部分组成，网络结构见图1。

Backbone被称为YOLOv5的主干特征提取网络，它由Focus结构和CSP结构组成。在输入图片经过CSPDarknet进行特征提取后，得到了一组被称为特征层的有效特征。接下来，从主干部分获取了三个特征层，并将它们用于进一步构建网络。

Neck部分采用了与YOLOv4相同的FPN + PANet结构。FPN用于加强特征提取，它能够融合不同尺度的特征信息。在上一步骤中获取的三个有效特征层会在这一部分进行特征融合。而PANet网络通过特征金字塔和特征融合两个模块的组合，有效地提升了目标检测性能。另外，SPP结构用于增强模型的感知能力和尺度不变性，而PAN结构则增强了多尺度特征的融合能力。

Yolo Head是YOLOv5的分类器和回归器，将上一步得到的特征图当作无数个特征点的集合，对特征点进行判定。非极大值抑制(nms)是为了避免对同一个物体进行多次检测，需要对重复的预测框进行过滤。

Figure 1. YOLOv5 network structure

图1. YOLOv5网络结构

3. 改进YOLOv5网络

本文以对矿石大小类别识别分割为目标，考虑到工业矿山生产线上工作量巨大，环境有限，目标检测定位需要具备较高的时效性，以便能顺利进行分拣工作，因此需要改进模型参数减少计算量，从而提升检测速度 [6] 。为了在速度与精度之间取得平衡，本文在主干特征提取网络引入轻量化的MobileNetv3，同时对加强特征提取网络PANet网络改进；将非最大抑制(nms)替换成DIou-NMS，可改善YOLOv5对重叠目标的识别。

3.1. MobileNetv3

MobileNetV3是一种轻量级网络，比起CSPDarknet该网络结构采用深度可分离卷积代替普通卷积，能有效降低模型参数量。并采用空间金字塔池化(spatial pyramid pooling, SPP)模块，对输出网络进行语义增强，利用空洞卷积掩码特征图增加感受野，如图2。

Figure 2. MobileNet specific bneck structure

图2. MobileNet特有的bneck结构

MobileNetv3在前两代的基础上推陈出新，综合了v1 [7] 的深度可分离卷积和v2 [8] 的具有线性瓶颈的逆残差结构，并引入了注意力机制，该机制由挤压模块(squeeze block)和激励模块(excitation block)组成。

MobileNetv3的结构如表1所示。

Table 1. Overall structure of MobileNetv3

表1. MobileNetv3整体结构

第一列表明了每层结构变化；第二列代表每层特征层即将经历的block结构；第三列则代表了逆残差结构升通道后的通道数，第四列是输入到逆残差结构时的通道数。SE表示每层是否引入了注意力机制；NL代表了激活函数的种类(h-swish和RELU)；s代表了每个block结构所用的步长。

取出其中三个尺度对应的有效特征层替换CSPDarknet53的原特征层。同时，为了达到实验要求的实时性和高精度，再将PANet网络3 × 3的卷积块改进为3 × 3深度可分离卷积和1 × 1的普通卷积，使参数量减少了24.95%。

3.2. NMS非极大抑制改进

原YOLOv5框架中使用CIoU，对于每个类别，根据置信度对所有候选框进行排序，将置信度最高的候选框作为起始框。遍历剩余的候选框，计算它们与起始框的重叠度(通常使用IOU，即交并比)。如果某个候选框与起始框的重叠度高于设定的阈值，则将该候选框剔除。

考虑到实际场地的矿石堆叠情况下，目标互相遮挡检测框的重叠面积较大，造成误检，使用DIoU和NMS结合方法改善此种情况，DIoU-NMS不仅考虑交并比的值，还考虑了预测边框和真实边界框之间中心距，如图3所示。

Figure 3. Schematic diagram of DIoU

图3. DIoU示意图

其中b、 $b^{gt}$ 是预测框和真实框的中心点，d是计算两个中心点间的欧氏距离，c是能同时覆盖两框的最小矩形的对角距离。

DIoU-NMS公式如下：

$s_i=\{\begin{array}{c}{s}_i,IoU-\frac{{\rho }^2\left(b,b^{gt}\right)}{c^2}\left(M,B_i\right)<\epsilon \\ 0,IoU-\frac{{\rho }^2\left(b,b^{gt}\right)}{c^2}(M,B_i)\ge \epsilon \end{array}$

其中 $s_i$ 用来表示分类分数， ${\rho }^2\left(b,b^{gt}\right)$ 代表的是计算两个中心点之间的欧氏距离， $M$ 表示预测分数最高的一个预测框， $B_i$ 判断预测框是否需要被移除， $\epsilon$ 用来表示NMS的阈值。

4. 实验结果与分析

4.1. 矿石数据集构建

图片样本采集于江西赣州钨矿，拍摄样本为选矿厂经过水洗工序后经由皮带运输带出的不同粒级原矿。随机挑选大小不一的原矿石置于实验平台上，放置好光源和相机，而后开始采集数据。预计样本1000张，最终仅拍摄200余张。

由于拍摄样本采集量不足可能导致模型过拟合，在进行训练前需对训练样本进行数据增强。数据增强常用重新定义尺寸大小，对给定的图像进行竖直/水平翻转、镜像翻转，图像裁剪、缩放比例，图片光感、亮度、曝光、对比度等。本文中样本经过裁剪处理后，将其进行了旋转(30˚、60˚、120˚)，缩放比例(0.5倍、0.75倍和1.25倍)，镜像/竖直/水平翻转变化，通过这些途径使数据集可达到目标张数。将数据按照9:1的比例划分为训练集和测试集，构建流程如图4所示。

Figure 4. Flow chart of ore data set construction

图4. 矿石数据集的构建流程图

4.2. 网络参数设置与模型训练

实验在Windows10系统下进行，GPU为NVIDIA GeForce RTX 3060 Laptop，显存6G，Pytorch为1.7.1，Python版本为3.7.12。训练时图片大小调整为640x640，训练轮数为200，初始学习率为0.01，动量因子为0.9。

改进的YOLOv5网络在数据集上的损失函数值变化曲线如图5所示。

Figure 5. Change curve of loss function value

图5. 损失函数值变化曲线

图5中epoch为模型训练时的迭代次数，损失值随着迭代次数的增加而下降。前十次的迭代中，损失值下降较快，10至100次迭代间，损失值下降平缓，参数接近最优值，150次后，训练集的损失值稳定在0.022左右，模型训练收敛且良好。

4.3. 评估指标

在分类任务中，典型的评估指标是准确性、精度、召回率和F1分数，其中准确性用于评估算法的整体性能。当数据集类别不平衡时，精度和召回率主要用于性能评估，而F1分数是精度和召回率的总和平均值。本文评估指标采用平均精度(AP)、平均精度均值(mAP)以及检测帧数(FPS)。衡量指标的定义如下：

$Accuracy=\frac{TP+TN}{TP+FN+FP+TN}$

$Precision=\frac{TP}{TP+FP}$

$Recall=\frac{TP}{TP+FN}$

其中TP (True Positive)表示真正例，TN (Ture Negative)表示真负例，FP (False Positive)表示假正例，FN (False Negative)表示假负例。

4.4. 实验分析

使用标定的矿石数据集在改进后的YOLOv5中进行训练，在相同的环境和参数设置下，与其他几种常用的目标检测算法进行对比。对比结果在下表2中。

Table 2. Comparison of the effect of common target detection algorithms

表2. 常用目标检测算法效果对比

从表中可以看出，相较于其他三种算法，本文算法在密集矿石堆叠情况下对矿石的检测表现效果更好，mAP值可以达到95.7%，在原始YOLOv5的基础上，帧数提高了62.5，检测单张图片的时间减少了4 ms，大小矿石类别的AP值分别提高了12.1%和0.7%，mAP值提高了6.4%。本文算法满足了矿山传送带检测的精度以及实时性。本文算法与YOLOv4、YOLOv5、YOLOv4-tiny算法进行对比的检测效果如图6所示。

原图 YOLOv4 YOLOv4-tiny YOLOv5 改进YOLOv5

Figure 6. Comparison of the effect of common target detection algorithms

图6. 常用目标检测算法效果对比图

从图6中可以看出，YOLOv4的检测效果表现较差；YOLOv4-tiny在检测效果上比YOLOv4略好一些，原始YOLOv5算法相较于前两种检测精度和定位更准，但仍然存在误判或者漏检的情况；与其他算法相比，本文的算法的目标检测效果表现突出，检测精度上升，误拣漏拣的现象基本没有，对重叠目标的鲁棒性显著提升。

5. 结论

本文在原有YOLOv5算法的基础上，分别从主干网络、NMS算法两个方面进行有效改进，改进后算法的mAP值达到了95.7%，帧数提高了59.99%，检测单张图片的时间缩短了4 ms，增强了YOLOv5网络模型对大块矿石的检测精度，缩短了检测时间。本文仅仅是对算法的优化改进，后期可以加入与机械手的协同，完善工业机械手实时分拣系统。

参考文献