一种基于改进YOLOv5模型的螺栓检测方法及系统

技术领域

本发明涉及计算机视觉技术领域，具体涉及一种基于改进YOLOv5模型的螺栓检测方法及系统。

背景技术

铁路轨枕螺栓是铁路轨道扣件的重要组成部分之一，由于其长期暴露在大气中，经受风吹雨淋，霜冻雪埋等循环性气候变化，容易出现氧化锈蚀。螺栓锈蚀轻则使螺栓锈死，影响轨距调整和轨道水平找正等养路作业，重则使螺栓根部腐蚀折断，影响行车安全。因此，为了保证运输安全，应有计划的对轨枕螺栓进行防锈蚀处理。近年来，在对轨枕螺栓进行长效防腐脂涂覆作业时，逐渐开始使用简单机具或智能装置进行螺栓自动识别和涂覆，然而，由于在识别中会出现遮挡、翻转以及平移等问题，

现有的钢轨螺栓识别方法往往使用计算机视觉技术，将图像处理技术用于钢轨螺栓的识别，可用于钢轨螺栓的自动涂覆，有效提高钢轨螺栓涂油保养的效率，减轻作业人员的工作负担。然而，由于钢轨螺栓图像中包含大量冗杂信息，往往还出现阴影、光照不均衡、复杂背景等情况，在识别中螺栓图像会出现遮挡、翻转以及平移等问题，造成误检和漏检的情况，的拓扑结构和形态差异较大，容易导致对其局部的误判和漏检。因此，如何快速、准确的识别钢轨螺栓，对提高铁路养路机械化水平，做好轨道零件日常性保养具有重要意义。

发明内容

为了解决上述问题，本发明提出了一种基于改进YOLOv5模型的钢轨螺栓检测方法，提高了识别效果，还减少了模型复杂度，提高了检测速度，能够快速的应对大规模的目标检测任务。

本发明解决其技术问题采取的技术方案是：

本发明实施例提供的一种基于改进YOLOv5模型的螺栓检测方法，包括以下步骤：

采集钢轨螺栓图像进行预处理，建立目标数据集；

构建结合CARAFE和NWD-IOU损失函数的改进YOLOv5模型；

利用目标数据集，对所述改进YOLOv5模型进行训练并评估；

使用评估完成的改进YOLOv5模型进行钢轨螺栓检测。

进一步的，采集钢轨螺栓图像进行预处理，建立目标数据集，包括：

使用LabelImg标注工具对采集的螺栓图像中螺栓位置进行标注，获得每张图像对应的xml文件；

对每张螺栓图像及其对应的xml格式文件进行数据增强操作；

将螺栓图像及其xml文件按照VOC格式进行存储，建立目标数据集。

进一步的，构建结合CARAFE和NWD-IOU损失函数的改进YOLOv5模型，包括：

建立基础YOLOv5模型；

将传统上采样模块替换为CARAFE上采样模块；

使用NWD-IOU损失函数作为模型的回归损失函数。

进一步的，CARAFE上采样模块进行上采样，包括：

输入形状为H×W×C的特征图，用1×1×1卷积将它的通道数压缩到H×W×C

设上采样倍率为σ，利用K

将所述上采样核利用Softmax进行归一化，使卷积核权重和为1。

进一步的，所述NWD-EIOU损失函数的公式，包括：

L

其中，IOU为预测边界框与真实边界框的重叠面积，IOU为使用IOU度量方法的阈值设定值，NWD为使用NWD度量方法的基础设定值。

进一步的，将所述目标数据集划分为训练集、验证集和测试集，其中训练集、验证集与测试集的比例为7:2:1。

进一步的，对改进YOLOv5模型进行评估，包括：

计算召回率Recall、精确率Precision和F1分数：

Recall＝TP/(TP+FN)

Precision＝TP/(TP+FP)

F1-Score＝2×(Precision×Recall)/(Precision+Recall)

其中，TP为将正类预测为正类的样本个数，FN为将正类预测为负类的样本个数，FP为将负类预测为正类的样本个数，TN为将负类预测为负类的样本个数。

本发明实施例提供的一种基于改进YOLOv5模型的螺栓检测系统，其特征在于，包括：

图像处理模块，用于采集钢轨螺栓图像进行预处理，建立目标数据集；

模型建立模块，用于构建结合CARAFE和NWD-IOU损失函数的YOLOv5模型；

模型训练模块，用于利用目标数据集，对所述改进YOLOv5模型进行训练并评估；

螺栓检测模块，用于使用评估完成的改进YOLOv5模型进行钢轨螺栓检测。

一种计算机设备，其特征在于，包括处理器、存储器和总线，所述存储器存储有所述处理器可执行的机器可读指令，当所述计算机设备运行时，所述处理器与所述存储器之间通过总线通信，所述处理器执行所述机器可读指令，以执行如权利要求1-7任一所述的一种基于改进YOLOv5模型的螺栓检测方法的步骤。

一种可读存储介质，其特征在于，所述可读存储介质上存储有计算机程序，所述计算机程序被处理器执行时实现如权利要求1-7任一所述的一种基于改进YOLOv5模型的螺栓检测方法的步骤。

本发明实施例的技术方案可以具有的有益效果如下：

(1)使用CARAFE上采样模块替换传统的上采样模块，具有更大的感受野，可以更好的利用周围的信息，与特征图的语义信息相关，基于输入内容进行上采样，不引入过多参数和计算量，使得模型进行了轻量化的改进；

(2)将NWD损失函数和IOU损失函数相结合，组成新型NWD-IOU损失函数，解决了使用IOU损失函数无法检测微小物体和使用NWD损失函数模型训练过程不收敛的问题；

(3)本申请基于YOLOv5算法结合轻量级上采样算子和NWD-IOU损失函数进行复杂背景下的螺栓检测，与传统算法相比，实际采集较少的螺栓图像进行训练，即能较好的处理阴影、气象环境、光照均衡、复杂背景和图像增强等难点问题。

附图说明

图1是根据一示例性实施例示出的一种基于改进YOLOv5模型的螺栓检测方法的流程图；

图2是根据一示例性实施例示出的改进YOLOv5模型的结构图；

图3是根据一示例性实施例示出的训练损失与改进YOLOv5模型评估结果；

图4是根据一示例性实施例示出的一种基于改进YOLOv5模型的螺栓检测系统的示意图。

具体实施方式

下面结合附图与实施例对本发明做进一步说明：

为能清楚说明本方案的技术特点，下面通过具体实施方式，并结合其附图，对本发明进行详细阐述。下文的公开提供了许多不同的实施例或例子用来实现本发明的不同结构。为了简化本发明的公开，下文中对特定例子的部件和设置进行描述。此外，本发明可以在不同例子中重复参考数字和/或字母。这种重复是为了简化和清楚的目的，其本身不指示所讨论各种实施例和/或设置之间的关系。应当注意，在附图中所图示的部件不一定按比例绘制。本发明省略了对公知组件和处理技术及工艺的描述以避免不必要地限制本发明。

如图1所述，本发明实施例提供的一种基于改进YOLOv5模型的螺栓检测方法，包括以下步骤：

步骤一：采集钢轨螺栓图像进行预处理，建立目标数据集。

采集钢轨螺栓图像，使用LabelImg标注工具对采集的螺栓图像中螺栓位置进行标注，获得每张图像对应的xml文件，所述xml格式文件用于标注对应图像中螺栓的位置。

对每张螺栓图像及其对应的xml格式文件进行数据增强操作，所述数据增强操作包括旋转、翻转变化、自适应亮度校正和裁剪等。

将螺栓图像及其xml文件按照VOC格式进行存储，建立目标数据集。

VOC的xml标注文件由两个级别的信息组成，第一个级别是图像整体信息，如所在文件夹，文件名，尺寸等。第二个级别的信息是object字段，object字段可以有多个，每一个表示图像中一个被标注的目标，因此每个object下面包含一些与目标相关的标注信息，比如目标类别，矩形框，特征点等。object字段下面可能还包含其他级别的信息，比如矩形框的标注就可以再往下延伸一级，因为一个矩形框包含四个数字(xmin，ymin，xmax，ymax)。

所述目标数据集划分为训练集、验证集和测试集，其中训练集、验证集与测试集的比例为7:2:1。

基于YOLOv5算法进行复杂背景下的螺栓检测需要大量带类别标签的螺栓图像作为训练集、验证集和测试集。但是截止目前，全球还没有公开的带类别标签、用于螺栓位置图像语义分割的数据合集，而本申请基于YOLOv5算法结合轻量级上采样算子和NWD-IOU损失函数进行复杂背景下的螺栓检测，通过实际采集得到5000个左右的螺栓图像并将其当做训练集、验证集进行训练，即能较好的处理阴影、气象环境、光照均衡、复杂背景和图像增强等难点问题。

步骤二：构建结合CARAFE和NWD-IOU损失函数的改进YOLOv5模型。

建立基础YOLOv5模型；将传统上采样模块替换为CARAFE上采样模块。基础YOLOv5模型结构包括：输入端、Backbone网络、Neck网络和输出端，YOLOv5的Head网络由三个不同的输出层组成，分别负责检测大中小尺度的目标，Backbone网络使用CSPDarknet53作为其主干网络，具有较强的特征提取能力和计算效率，Neck网络使用FPN网络，能够在不同的特征图层次上进行检测，可以提高目标检测的性能，融合来自不同特征图层次的信息。本申请对Backbone网络中的上采样模块进行改进，如图2所示，将传统上采样模块修改为CARAFE上采样模块，传统的上采样模块仅通过像素点的空间位置来决定上采样核，并没有利用到特征图的语义信息，可以看作是一种“均匀”的上采样，而且感知域通常都很小。对比于传统的上采样模块，CARAFE上采样模块具有较大的感受野，可以更好的利用周围的信息，与特征图的语义信息相关，基于输入内容进行上采样，并且不引入过多的参数和计算量，对模型进行了轻量化改进。

CARAFE分为两个主要模块，分别是上采样核预测模块和特征重组模块。当上采样倍率为σ，给定一个形状为H(图片高度)×W(图片长度)×C(图片通道数)的输入特征图，首先利用上采样核预测模块预测上采样核，然后利用特征重组模块完成上采样，得到形状为σH×σW×C的输出特征图。CARAFE上采样模块进行上采样的步骤，具体为：

上采样核预测模块进行特征图通道压缩、上采样核预测和上采样核归一化。输入形状为H×W×C的特征图，用1×1×1卷积将它的通道数压缩到H×W×C

特征重组模块对于输出特征图中的每个位置，将其映射回输入特征图，取出以之为中心的K

使用NWD-IOU损失函数作为模型的回归损失函数，所述NWD-EIOU损失函数的公式，包括：

L

其中，IOU为预测边界框与真实边界框的重叠面积，IOU为使用IOU度量方法的阈值设定值，NWD为使用NWD度量方法的基础设定值。

本发明在损失函数的改进上提出了使用NWS-IOU算法，其中NWD是专门提出用于小物体的目标检测。因为微小物体仅包含几个像素大小，由于缺乏外观信息，最先进的检测器无法在微小物体上产生令人满意的结果(IOU)，并且在用于基于锚点的检测器中时会大大降低检测性能。现有方法往往使用使用Wasserstein距离对IOU进行替换进行微小物体检测：首先将边界框建模为2D高斯分布，然后提出一个称为归一化Wasserstein距离(NWD)的新度量，以通过它们对应的高斯分布来计算它们之间的相似性。与IOU相比，NWD在微小物体检测方面具有尺度不变性、对位置偏差的平滑性，并且能够测量不重叠或相互包容的包围盒之间的相似性，因此在微小物体检测方面有比IOU更好的特性。但是在YOLOv5只使用NWD损失函数的情况下，模型在训练过程中，会出现收敛速度慢或者不收敛的情况，针对此类情况，本发明将NWD损失函数和IOU损失函数相结合，从而组成为新型的NWD-IOU损失函数，用于解决使用NWD损失函数模型训练过程不收敛的情况，使得函数在收敛快的同时具有更好的识别效果，加入了IOU设定值(IOUr)，通过预先对IOUr进行设定，可以针对于图片中的目标大小从而更改损失函数的设定。

步骤三：利用目标数据集，对所述改进YOLOv5模型进行训练并评估。

本发明模型训练过程中采用常规的训练方式，其中超参数的设置如表1所示：

表1超参数设置

具体的参数解释如下：

Epochs(批次)：是指将整个数据集迭代一遍的过程。在一个Epoch中，模型会对整个数据集进行一次前向传播和反向传播，更新所有的参数。

Batch-size(批尺寸)：表示单次传递给程序用以训练的数据(样本)个数。

Img-size(图片大小)：表示原始输入图片输入到模型中的预设大小。

Learning-rate(学习率)：学习率。

Momentum(动态冲量)：优化算法是对梯度下降法的一种优化，它在原理上模拟了物理学中的动量，已成为目前非常流行的深度学习优化算法之一。

Weight_decay：权值衰减又为L2正则化。

Device：训练使用的处理方式，0为gpu处理。

本实施例对改进YOLOv5模型进行评估，使用的评估标准为以下数值：

(1)召回率：又称查全率，表示在全部正值样本总数中，预测到的真实正例的比重，以反映漏检情况。

Recall＝TP/(TP+FN)

(2)精确度：也称为查准率，指的是在已被预测的正例中所占的百分比，用来度量一个正例在一个预测中被预测成真的可能性。

Precision＝TP/(TP+FP)

(3)F1-Score：为精确率和召回率的调和平均值，同时兼顾了分类模型的精确率和召回率。

F1-Score＝2×(Precision×Recall)/(Precision+Recall)

其中，TP为将正类预测为正类的个数，FN为将正类预测为负类的个数，FP为将负类预测为正类的个数，TN为将负类预测为负类的个数。

本申请改进YOLOv5模型训练损失函数和评估值变化如图3所示，第一行展示的分别是边界框训练集损失、目标训练集损失、精确度、召回率，第二行展示的是边界框验证集、目标验证集损失、平均精准度在IoU＝0.5时的数值、平均精度在IoU＝0.5-0.95总和数值，从表中可以看出第一行，在训练过程中，边界框损失和分类损失逐渐降低并在迭代20轮后趋于平缓，而精确度和召回率在20轮之后准确率可达到0.9以上，图中第二行展示的边界框验证损失、分类验证损失也是在迭代20轮之后趋于平缓，而在IoU＝0.5的时候平均准确率也可达到0.9以上。

步骤四：使用评估完成的改进YOLOv5模型进行钢轨螺栓检测。

如图4所示，本发明实施例提供的一种基于改进YOLOv5模型的螺栓检测系统，包括：

图像处理模块，用于采集钢轨螺栓图像进行预处理，建立目标数据集；

模型建立模块，用于构建结合CARAFE和NWD-IOU损失函数的YOLOv5模型；

模型训练模块，用于利用目标数据集，对所述改进YOLOv5模型进行训练并评估；

螺栓检测模块，用于使用评估完成的改进YOLOv5模型进行钢轨螺栓检测。

本发明实施例还提供了一种计算机设备，包括处理器、存储器和总线，所述存储器存储有所述处理器可执行的机器可读指令，当所述计算机设备运行时，所述处理器与所述存储器之间通过总线通信，所述处理器执行所述机器可读指令，以执行如上述任意所述的一种基于改进YOLOv5模型的螺栓检测方法的步骤。

本发明实施例还提供了一种存储介质，该存储介质上存储有计算机程序，该计算机程序被处理器运行时执行如上述任意所述的一种基于改进YOLOv5模型的螺栓检测方法的步骤。

本领域内的技术人员应明白，本申请的实施例可提供为方法、系统、或计算机程序产品。因此，本申请可采用完全硬件实施例、完全软件实施例、或结合软件和硬件方面的实施例的形式。而且，本申请可采用在一个或多个其中包含有计算机可用程序代码的计算机可用存储介质(包括但不限于磁盘存储器、CD-ROM、光学存储器等)上实施的计算机程序产品的形式。

本申请是参照根据本申请实施例的方法、设备(系统)、和计算机程序产品的流程图和/或方框图来描述的。应理解可由计算机程序指令实现流程图和/或方框图中的每一流程和/或方框、以及流程图和/或方框图中的流程和/或方框的结合。可提供这些计算机程序指令到通用计算机、专用计算机、嵌入式处理机或其他可编程数据处理设备的处理器以产生一个机器，使得通过计算机或其他可编程数据处理设备的处理器执行的指令产生用于实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能的装置。

这些计算机程序指令也可存储在能引导计算机或其他可编程数据处理设备以特定方式工作的计算机可读存储器中，使得存储在该计算机可读存储器中的指令产生包括指令装置的制造品，该指令装置实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能。

这些计算机程序指令也可装载到计算机或其他可编程数据处理设备上，使得在计算机或其他可编程设备上执行一系列操作步骤以产生计算机实现的处理，从而在计算机或其他可编程设备上执行的指令提供用于实现在流程图一个流程或多个流程和/或方框图一个方框或多个方框中指定的功能的步骤。

最后应当说明的是：以上实施例仅用以说明本发明的技术方案而非对其限制，尽管参照上述实施例对本发明进行了详细的说明，所属领域的普通技术人员应当理解：依然可以对本发明的具体实施方式进行修改或者等同替换，而未脱离本发明精神和范围的任何修改或者等同替换，其均应涵盖在本发明的权利要求保护范围之内。