基于改进算法YOLOv5+的混凝土轨枕裂纹检测

令雅莉，杨桂芹，张又元，王小鹏

(兰州交通大学电子与信息工程学院,兰州 730070)

引言

铁路运输安全直接关系到人民的生命安全与社会的经济发展。裂纹作为主要的轨枕损伤,当细小裂纹的出现未及时发现时,细小裂纹会扩大成裂缝,继而发展为轨枕失效,对行车安全造成威胁。因而,定期对轨枕进行检测,在裂纹较小的时候,及时发现并解决,能够降低维修成本,提高铁路的基础设施本体设备安全。当前,铁路轨枕主要以人工巡检配合机器辅助的检测方法。在检测过程中,这种方式主要存在人力耗费大、长时间巡检容易导致巡检人员疲劳,造成错检和漏检、裂纹的信息需要手动输入,无法实现轨枕的自动化检测,速度太慢,效率低下。基于以上原因,非常有必要实现混凝土轨枕裂纹自动化检测。

目前,实现自动化检测混凝土轨枕裂纹的研究方法主要有两大类,第一类是以RCNN[1]、Faster-RCNN[2]等为代表的双阶段检测算法。该算法一般将目标检测分为两步,第一步主要是在图片上获得候选区域(滑动窗口),从而实现候选区域特征向量的提取。第二步是分类和位置预测候选区域。双阶段检测算法虽然可以提高检测准确率,但由于扩大算法模型导致计算量增加,检测速度减慢,无法满足实时性检测的需求。第二类是以SSD[3]、YOLO[4-6]为代表的单阶段检测算法,该类算法省略了候选区域生成阶段,径直得到目标分类和位置信息,极大提升了目标检测的速度。虽然这两类算法都有相应的优点,但检测速度慢与检测精度低的缺点导致其不适用于铁路混凝土轨枕裂纹检测。

考虑到双阶段算法检测速度难以提升。因此,本文选择单阶段算法YOLOv5为基础,提出了一种基于改进YOLOv5m[7]的铁路轨枕裂纹检测方法。该方法在骨干网中嵌入SE[8]注意力机制模块对主干网络进行调整,获取更多浅层特征从而得到更多的图片纹理信息。文献[9]通过热力图证明最大池化很容易使裂纹信息丢失,因此,将SPP模块中的最大池化层更换为平均池化层来保留更多裂纹信息。设计出一种新的多尺度检测结构,在保留速度的前提下,减少微小裂纹的漏检情况,提高裂纹的检测精度。考虑裂纹的尺度变化范围过大会影响单阶段算法的检测精度,所以采用分治标签的方法人为降低裂纹尺度变化范围,减少微小裂纹的漏检,从而提高检测精度。

1 YOLOv5m网络模型介绍

YOLOv5是一种单阶段目标检测算法,与YOLOv4[6]对比来看,其具有均值权重文件更小,训练时间和推理速度更短的优势,YOLOv5一共有4个版本, 按照网络深度大小和特征图宽度大小分为YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x。基于对识别精度、检测速度以及权重文件大小等因素的综合考虑,选择检测速度较快、识别精度较高的YOLOv5m进行研究与改进,其权重文件也相对较小,适合部署在嵌入式设备中,具有很高的实用性。YOLOv5m模型主要由Input、Backbone和Head这3个部分组成。如图1所示。

图1 YOLOv5网络结构Fig.1 YOLOv5 Network structure

Input主要包括自适应缩放图片、Mosaic数据增强以及自适应锚框计算。因为不同的图片有不同的尺寸,所以采用自适应缩放图片可以将原始图片统一缩放到一个符合模型需求的标准尺寸,然后送入检测模型中训练,本文的输入图片尺寸为640×640像素;Mosaic数据增强是利用了4张图片,对这四张图片进行拼接后得到一张新的图片,通过增加小目标的方式极大提升小物体的识别能力,使数据更加充足;自适应锚框计算将开始设定的锚框作为基础从而输出预测框,接着和真实框进行对比分析,采用K均值和遗传算法对自定义数据集实行计算,算出其中的差别后进行反向变换,最后一直轮换参数得到最恰当的锚框值。

2 YOLOv5优化改进

2.1 YOLOv5在轨枕裂纹检测方面的不足

虽然YOLOv5模型具有识别速度快、检测精度高的特点,但利用既有的YOLOv5模型对混凝土轨枕裂纹进行专有化检测,存在以下不足。

(1)混凝土轨枕中存在大量微小裂纹,通常在提取的过程中,小目标物体凭借浅层特征更多一点,因此极易遗失小目标物体信息。

(2)最大池化容易导致裂纹特征信息丢失,也是导致微小裂纹漏检的主要原因。

(3)当输入图像尺寸为640×640像素时,在原始模型中,检测分别使用8倍、16倍、32倍下采样相应得到尺寸大小分别为80×80、40×40、20×20像素的检测层,分别可以检测大小在8×8、16×16、32×32像素以上的目标,但是对于轨枕裂纹检测而言,其存在大量的微小裂纹,而且采用分治标签的方法也会使裂纹分为若干段的小裂纹去检测。因此,裂纹目标物都变为小尺寸目标。在检测小于8×8像素的微小裂纹时会变为若干段的小裂纹去检测。因此,裂纹目标物都变为小尺寸目标。在检测小于8×8像素的微小裂纹时会变得困难,过大的检测层也变得多余。

2.2 网络模型改进

2.2.1 骨干网改进

由于混凝土轨枕的表面情况相比一些典型应用场景更为复杂,且大量早期裂纹的边缘相较一般物体与背景更难区别。原始YOLOv5模型在进行卷积下采样时容易使小目标的特征信息丢失,因此在骨干网络中引入SE注意力机制来强调更为重要的裂纹信息,抑制不太重要的背景信息,尤其对于小裂纹和密集裂纹来说,使模型更加精准地定位和识别裂纹区域。

SENet是Huetal于2017年所提出的融合通道注意力与空间注意力所得到的网络结构,具体结构如图2所示,其中W,W′,H,H′为特征图宽高;C,C′为通道数;Fsp为压缩操作,即进行全局平均池化;Fex为激励操作,即为减少通道个数从而降低计算量;Fscale是通道权重相乘,输入特征图大小为W′×H′×C′,最后输出特征图大小为W×H×C。

图2 SE(Squeeze-And-Excitation)模块Fig.2 Squeeze-And-Excitation module

从图2中可以看出,SE模块主要包括Squeeze和Excitation两部分,并通过这两部分来处理全局信息。卷积会得到Feature map,首先对Feature map进行处理,处理后可以得到一个与通道数量相同的一维向量,这个一维向量可以成为每个通道的评价指标,接着将指标分配到对应的通道上,得到最终结果。为使网络能够更好地拟合通道之间的相关性,获取更重要的通道特征的信息,在YOLOv5主干网中加入SE-Net通道注意力模块,即在主干网络结尾C3模块之后,添加SENet模块,并与之后的Head层相连。通过建立特征映射关系,使网络充分利用得到的全局信息,将裂纹目标的通道特征信息赋予更高的权重。从而更好地拟合裂纹目标通道间的相关特征信息,抑制无用的信息,最终会使模型更专注于检测裂纹这一指定类别。其模型构建流程如图3所示。

图3 SE-ResNet流程Fig.3 SE ResNet process

2.2.2 基于平均池化的SPP模块

原始的SPP模块主要通过5×5、9×9、13×13像素的最大池化和填充操作与未池化数据进行concat拼接来增大感受野,提取出最重要的多方面特征。最大池化在典型场景下会有更好的效果,但对于裂纹目标来说,却会使得特征信息丢失。因此,本文将SPP模块中的最大池化替换为平均池化,在具备增大感受野和提取多方面特征的优势下,来保留更多裂纹特征。改进前后分别如图4、图5所示。

图4 基于最大池化的SPP模块Fig.4 SPP module based on maximum pooling

图5 基于平均池化的SPP模块Fig.5 SPP module based on average pooling

2.2.3 检测结构改进

YOLOv5采用了3种尺度对不同大小的目标进行检测。分别通过8倍、16倍以及32倍下采样特征图对小型、中型以及大型目标进行检测[10-11]。原有多尺度检测结构对细小裂纹目标容易造成漏检,且分治标签的方法缩小了裂纹尺度变化范围,使大段裂纹目标变为小段裂纹目标检测。因此,针对该问题,改进了原有的检测结构。

首先,删除原始模型中32倍下采样所得到的尺寸为20×20像素较大物体的检测尺度;增加一个大小为160×160像素的特征图。具体操作如下:在增加了SE注意力机制后的网络结构中第18层之后继续对特征图进行上采样操作,持续增加特征图,在第21层时,将160×160像素的特征图与骨干网络中第2层特征图进行张量拼接融合,从而通过新的特征融合获取到更大的特征融合,对小目标进行检测。在第22层增加小目标检测层。upsample为上采样模块,经过上采样操作会扩大特征图。改进后的网络结构如图6所示。

图6 改进后的YOLOv5网络结构Fig.6 Improved YOLOv5 network structure

3 实验与分析

3.1 数据集介绍

本文采用文献[9]中采集的数据集,从中选取2 086张高质量的裂纹图片整理为本次实验的数据集,数据集样例如图7所示。按照8∶2的比例将其划分为训练集和测试集,其中训练集有1 669张,测试集有417张。

图7 部分数据集样例Fig.7 Example of part dataset

常规标注方式有3点不足:(1)会使得目标框尺度变化过大,难以收敛;(2)容易使微小裂纹漏检;(3)通过预选锚框与真实标签的交并比计算容易将正确的标注框归类于负样本。因此,针对以上常规标注方式的不足之处,本文采用分治标签的方式并使用Labelimg对裂纹进行标注,具体标注方式如图8所示。

图8 数据集标注Fig.8 Dataset annotation

分治标签使用了适合裂纹尺度的目标框,将裂纹被漏检的风险降到最低。图9为数据集的可视化结果,其中图9(a)为数据集目标类别的分布与目标的数量,图9(b)为数据集目标中心点位置分布,图9(c)为数据集目标大小分布,横纵坐标分别代表目标的宽和长。

图9 数据集可视化结果Fig.9 Dataset visualization results

(1)

3.2 目标框回归损失函数

常用的目标框回归损失函数有GIoU[12]或CIoU[13]。GIoU能同时关注到真实框与预测框的重叠面积和其他的非重叠区域,所以GIoU与原始的IoU相比可以更好地反映两者之间的重合度。但考虑GIoU始终只能注意到真实框与预测框之间重叠率这一个因素,无法更好地描述目标框的回归问题,采用CIoU_loss作为全文检测任务的回归损失函数。CIoU_loss[14-15]可以关注预测框与目标框之间的重叠面积和中心点距离,同时也考虑预测框与目标框的长宽比,使模型回归效果更好,CIoU_Loss计算公式如式(1)所示。

式中,ρ2(b,bgt)为预测框与实际框两个中心点之间的距离;c为最小闭合区域的对角线距离;α为权重函数;υ为用来度量检测框和真实框之间的比例的相似性。α和υ分别如式(2)、式(3)所示。

(2)

(3)

式中,hgt为目标框的长;wgt为目标框的宽;h为预测框的长;w为预测框的宽。

3.3 实验环境与模型训练

本文实验的操作系统为Windows10,运行内存为12GB,CPU型号为I5,GPU为GeForce RTX 2080 Ti,显存大小为11GB,实验基于Pytorch1.7.1深度学习框架,编程语言为Python3.8,版本为CUDA 11.2。

训练采用SGD优化算法,初始学习率为0.01,采用余弦退火策略动态降低学习率,动量因子为0.937,权重衰减系数为0.0005,Batch size为16,Epochs为100,将前3次Epoch做为模型预热。

3.4 评价指标

本文采用通用评价指标,即精确率(Precision)、召回率(Recall)、平均精度均值(mean Average Precision,mAP)以及帧率(Frames Per Second,FPS)对网络模型进行评价[16]。

通常在二分类问题中,常用混淆矩阵(Confusion Matrix)来表示样本预测值的正负与样本真实值的正负之间的关系[17],图10表示一个混淆矩阵。

图10 混淆矩阵Fig.10 Confusion matrix

召回率与精确率用公式表示[18],分别如式(4)、式(5)所示。

(4)

(5)

式中,FPrecision为预测裂纹中正确预测的数量占总预测裂纹数量的比例;FRecall为预测裂纹中正确预测裂纹数量占总正确预测数量的比例[19]。FTP表示将正样本预测为正;FFP表示将负样本预测为正;FTN表示将正样本预测为负;FFN表示将负样本预测为负。

FAP指PR(Precision-Recall)曲线下面积,FmAP指每个类别FAP的均值[20],分别如式(6)、式(7)所示。

(6)

(7)

式中,p(r)为不同查准率r下对应的查全率p;APi为第i类的检测准确率;N为分类数。

检测速度通常用FPS来衡量,表示检测模型每秒处理图片的数量,FPS值越高,网络模型处理图像的速度就越快[13]。

3.5 实验结果对比与分析

本文共有2组实验。先将改进后的算法YOLOv5+与既有的YOLOv5算法对混凝土轨枕裂纹检测进行详细对比,然后又与Faster R-CNN、YOLOv3模型进行了性能对比。

改进算法在训练100个迭代周期过程中损失变化情况如图11所示,由图可知,前期的损失值以较快的速度降低,但由于迭代次数逐步的增加,Loss值逐渐下降,处于一个平稳的状态,改进后的网络模型收敛速度很快。

图11 Loss变化曲线Fig.11 Loss variation curve

本文算法与原始YOLOv5算法平均精度均值、精确率和召回率的变化曲线对比分别如图12～图14所示,详细数据对比如表1所示。在训练过程中,改进后的YOLOv5模型与原始模型都采用同一个数据集和相同的参数设置,从图中和表中可以看出,改进后P(r)的模型在平均精度均值、精确率和召回率的对比中,均优于原始模型,且损失值变化更加平稳。改进后的模型检测速率比原始模型更快的同时模型参数量也更少,YOLOv5中为21 037 638,本文算法为17 600 742。

表1 模型改进前后参数详细对比 %Tab.1 Detailed comparison of parameters before and after model improvement

图12 改进前后mAP对比Fig.12 Comparison of mAP before and after improvement

图13 改进前后Precision对比Fig.13 Comparison of Precision before and after improvement

图14 改进前后Recall对比Fig.14 Comparison of Recall before and after improvement

为能够进一步对比本文改进算法的检测性能,将本文改进后使用的YOLOv5+算法与普遍使用的目标检测模型Faster R-CNN、YOLOv3算法进行实验对比分析,采用Precision、Recall和mAP三项指标分别进行评估与对比,对比实验结果如表2所示。

表2 不同算法性能对比结果Tab.2 Performance comparison results of different algorithms

由表1、表2可知,本文算法Precision比Faster R-CNN、YOLOv3与YOLOv5算法分别高出15.8%、7.98%、6.5%,Recall与这3种算法相比差不多,但从mAP来看,本文算法均高于其他算法,且比典型双阶段算法Faster R-CNN高出近20%,比YOLOv3高出了近10%,比YOLOv5高出了8%,而且从表2可以看出FPS有所提高。

经综合比较,在Recall相差不多的情况下,本文无论从Precision还是mAP,都有显著的优越性,而且实时性比YOLOv5更高,故本文算法相较于其他算法更适合混凝土轨枕裂纹检测任务。

3.6 检测结果对比与分析

为能够更进一步考证改进模型的切实性,从测试集中选择一些图片进行测试,图15为原始YOLOv5模型与改进后模型在不同条件下的检测结果对比。

图15 改进前后检测效果对比Fig.15 Comparison of detection effects before and after improvement

由图15(a)可知,在光线较亮、裂纹特征较明显的条件下,YOLOv5算法与本文改进算法均取得良好的检测效果,但YOLOv5算法仍存在小部分裂纹漏检的情况。由图15(b)和图15(c)可知,在小段细微裂纹的检测中,YOLOv5算法对微小裂纹的漏检情况较为明显,而本文改进算法对特征不明显的裂纹检测效果会好很多。由图15(d)可知,原算法在对破碎裂纹检测时,存在大量漏检,而改进后算法对破碎裂纹的检测效果良好。由图15(e)可知,在光线暗的条件下,原算法未检测出单个微小裂纹,而改进后算法成功检测出了轨枕所存在的微小裂纹。由图15(f)可知,在轨枕表面条件相对比较繁琐的情况下,在进行裂纹检测时,原算法受到一定干扰,存在多处漏检情况,改进后算法检测效果优异。

4 结语

在已有的YOLOv5算法对混凝土轨枕裂纹检测效率不高,容易造成细小裂纹漏检的基础上,基于改进算法YOLOv5+的混凝土轨枕裂纹检测针对微小裂纹检测效果更好。现阶段采用的深度学习目标检测算法对轨枕裂纹检测效果不理想,尤其是检测微小裂纹时容易出现漏检现象,因此,本文提出了一种改进算法,该改进算法主要基于YOLOv5模型,改进部分主要是通过在YOLOv5的主干提取网络中加入SELayer注意力机制模块、更换最大池化层为平均池化层、构建新的多尺度检测结构来提高对轨枕裂纹检测的能力并减少对微小裂纹漏检的情况。并选用分治标签的方法对混凝土轨枕裂纹进行标注,从而缩小裂纹尺度变化范围,加快模型收敛,增大裂纹的实际占比并提高对微小裂纹的检测效果。实验与检测结果表明,本文改进方法无论在检测精度还是速度上都优于YOLOv5原算法,且参数量更少,较好地满足实时性检测要求,总体性能也高于Faster R-CNN、YOLOv3等目标检测模型。