改进YOLOv5s的明渠漂浮垃圾实时检测方法

石露露,廖光忠

(1.武汉科技大学 计算机科学与技术学院,湖北 武汉 430065;2.武汉科技大学 智能信息处理与实时工业系统湖北省重点实验室,湖北 武汉 430065)

0 引 言

南水北调工程是目前解决国内北方地区水资源严重短缺问题的特大型基础设施项目。确保调水水质是其成败的关键,但是这项工程的沿途都是明渠,水面难免会有漂浮的垃圾,传统处理明渠漂浮垃圾的方法主要是依靠人力,这种方式不仅耗时耗力,而且无法做到实时监测。随着无人机航拍和无人船航拍技术的兴起,人工巡检逐渐被取代,如何检测航拍图像上明渠漂浮垃圾成为了亟待解决的问题。

近年来,随着深度学习的快速发展,基于深度学习的目标检测算法成为了检测航拍图像上明渠漂浮垃圾的主流方法。基于深度学习的目标检测算法主要分为双阶段目标检测和单阶段目标检测。双阶段目标检测算法虽然检测的准确率高,但是速度较慢,无法实现实时检测。单阶段目标检测算法是直接对输入的图像进行卷积特征提取,不需要使用候选框,可以极大地提高检测速度。因此,越来越多的学者开始将单阶段目标检测算法应用于河道漂浮物检测领域。Li等[1]基于改进的YOLOv3[2]算法实现水面漂浮垃圾检测,将三尺度检测转换成两尺度检测以提高检测速度,并通过调整先前的锚框提高检测精度。唐小敏等[3]基于ResNet-101特征提取网络的SSD网络模型对河道漂浮物进行目标检测,有效地提高了网络对不同尺度漂浮物的识别精度。李德鑫等[4]对YOLOv5s算法进行改进实现河道漂浮垃圾分类检测,降低各类别目标检测精度之间的差异。以上单阶段目标检测算法虽然在精度上有了一定的提升,但是小目标检测精度仍不高。针对小目标检测问题,国内外学者也提出了许多算法。Benjumea等[5]提出了YOLO-z,即通过BiFPN[6](Bidirectional Feature Pyramid Network,加权双向特征融合网络)代替YOLOv5中的PANet[7](Path Aggregation Network,路径聚合网络),扩大Neck层等一系列方法促进中浅层的特征融合,从而提升小目标的检测性能。Zhu等[8]提出了TPH-YOLOv5,在YOLOv5的基础上增加一个预测头用于检测不同尺度的目标,还通过集成CBAM[9](Convolutional Block Attention Module)注意力机制来寻找对象密集场景下的注意力区域,有效地提升了小目标的检测性能。目前现有的大部分小目标检测算法应用领域比较局限,一般只在大型公共数据集上进行验证,而现实场景错综复杂,现有算法仍然存在检测精度低、难以准确检测到小目标、速度不满足日常需要、漏检和误检等问题。

在明渠漂浮垃圾检测的场景中,航拍图像中的漂浮垃圾大部分尺寸都很小,占用的图像区域都比较小,导致这些漂浮垃圾缺乏外观信息,很难与背景区分开来,从而增加了准确检测明渠漂浮垃圾的难度。而且水面的强光反射可能会导致图像过度曝光,使水面漂浮垃圾不可见,增加了明渠漂浮垃圾检测的漏检率。此外,水面的倒影会对漂浮垃圾检测造成干扰,容易产生误检。

为了解决以上问题,该文采用的方法如下:

(1)通过随机剪裁、水平翻转、加入高斯噪声、改变对比度等图像处理操作进行数据增强来扩充数据集;

(2)借鉴BiFPN的结构对YOLOv5s原有的Neck 部分进行改进,修改特征融合的过程,提高小目标检测检测精度和检测速度;

(3)进行模型优化,在Neck和Head部分之间添加3个改进的CBAM注意力机制,增强信息的特征提取和定位;同时增大网络的输入分辨率,使图像信息更加丰富和精确。

1 YOLOv5s基本模型

2020年6月,Ultralytics团队Glenn等人提出了YOLOv5模型,虽然在性能方面与YOLOv4[10]不分伯仲,但是推理速度有非常大的提升,可以满足实时检测的要求。YOLOv5包括YOLOv5s、YOLOv5m、YOLOv5l和YOLOv5x四种网络结构[11]。该文主要是基于YOLOv5s进行改进的。为了可以在提高检测精度的同时,也能保证检测速率,以实现实时检测,YOLOv5s模型主要由Input、Backbone、Neck、Head四部分组成。

Input部分主要包括自适应计算锚框、Mosaic数据增强和自适应缩放图片三个部分。Backbone部分包括CSPNet[12](Cross Stage Partial Network,跨阶段局部网络)模块和SPPF模块,这部分的作用是为了处理输入图像,将图像转换成特征。YOLOv5s 6.0为了避免Focus切片下采样会丢失部分图片原始信息,将Focus换成了普通的卷积块。此外,YOLOv5s 6.0还用SPPF替代了原先的SPP(Spatial Pyramid Pooling,空间金字塔池化),SPP的3个多尺度最大池化层是并联连接的,而SPPF在此基础上将3个多尺度最大池化层串联起来,保留SPP优势的同时,也可以减少图片信息的丢失。CSPNet可以缓解梯度消失的问题,避免梯度的重复计算,从而减少计算量,提升网络的计算能力以及减少内存的损耗。Neck部分包括FPN[13](Feature Pyramid Networks,特征金字塔)和PANet,这部分主要是为了充分地提取特征。Head部分主要是对图像的特征进行预测,然后生成边界框、预测目标的种类。同时采用了三种损失函数分别计算分类损失、定位损失和置信度损失,并通过NMS(Non-Maximum Suppression,非极大值抑制)提高网络预测的准确度。

2 改进的YOLOv5s模型

针对航拍明渠漂浮垃圾图像存在的问题,主要从数据增强、多尺度特征融合、引入注意力机制、增加网络的输入尺度这四个方面对YOLOv5s模型进行改进和优化,改进的YOLOv5s的网络结构如图1所示。

图1 改进的YOLOv5s的网络结构

2.1 数据增强

由于目前没有用于河道漂浮检测的大型公共数据集,该文使用了来自FloW[14]的图像子数据集FloW-Img子集和在线收集的2 000张数据。FloW数据集是全球第一个无人船视角的水面漂浮垃圾检测数据集。该数据集是使用无人船在真实场景下进行拍摄收集的,但是FloW-Img数据集中的漂浮垃圾主要包括不同外观的塑料瓶、饮料瓶、草和落叶,数据种类较少。于是在FloW-Img数据集中选取1 000张优质图片,剩余1 000张则是在线收集合适的水面漂浮垃圾图像,确保数据集具有多样性。最后数据集包括落叶、塑料袋、瓶子、牛奶盒、塑料垃圾、草、树枝、球8个类别,如图2所示。

图2 水面漂浮垃圾类别

为了避免数据集太小而出现过拟合问题,将通过局部剪裁、随机翻转、加入高斯噪声、改变对比度等方式进行图像处理操作,让每张原始图像生成5张新图像,总共得到10 000张,但是经过数据增强之后,有的图片质量变化较为严重,需要将质量较差的图像进行人工筛选和剔除,最终数据集为8 400张。在图像处理时随机使用1到3种方式进行变换,可以增加数据集中图片的数量,提高图片信息的多样性,使模型在训练时可以学习到更加全面的水面漂浮垃圾特征。之后,通过数据标注工具Labelimg对所有图片进行标注,并将数据集所有图像按照7∶2∶1的比例划分成三个数据集,即训练数据集、验证数据集和测试数据集,图片数量分别为5 880张、1 680张、840张。最后,将VOC格式的数据标注文件转换为YOLO格式。

2.2 多尺度特征融合

在卷积神经网络的多层结构中,浅层网络更关注局部细节信息,如图像边缘、轮廓、纹理等,有利于对目标进行定位;深层网络则侧重于提取抽象信息,有利于对目标进行分类,但是细节信息比较缺乏。因此,将浅层网络特征与深层网络特征进行融合可以提高目标检测的性能。通常,构造特征金字塔FPN是一种直接有效的融合方式,即建立一条自顶向下的通路进行特征融合,然后用融合后具有高语义信息的特征层进行预测,但是这种结构受单向信息流的限制。为了解决这个问题,PANet结构应运而生,PANet结构是在FPN结构的基础上再添加一条自底向上的通路,这条通路主要用于将底层的位置信息也传到预测特征层中,使预测特征层既有丰富的语义信息也有精确的位置信息,不仅可以提高目标检测的精度,而且有利于小目标检测,如图3(a)所示。

(a)PANet结构 (b)Bi-PANet结构

虽然PANet可以提高目标检测的精度,但是需要更多的参数和更大的计算量。于是,借鉴BiFPN加权双向特征融合的思想将原先PANet的特征融合过程进行修改,如图3(b)所示,称为Bi-PANet。首先在PANet的基础上去掉P5、P3中没有进行特征融合的不必要连接,使融合更多特征的同时消耗较少的成本,减少计算量;然后将图3(b)虚线方框里的双向路径作为一个特征网络层,重复叠加相同的特征网络层, 以实现更高层次的特征融合;最后,给每个输入特征增加一个额外的权重,让网络了解它们的重要性。其中,采用的加权特征融合方法是快速归一化融合,每个归一化权重的值都是在0到1之间,可以提高计算速度。快速归一化融合的公式如下所示:

(1)

其中,ωi是可学习的权重,用激活函数ReLu来确保每个ωi≥0,该权重的值可由网络训练得到,Ii是指输入的特征,O是指输出特征,ε=0.000 1是避免数值不稳定的一个小值。

Bi-PANet将主干网络的P3、P4、P5作为特征的输入。以节点P4为例,其形成2个融合特征的过程如下所示:

(2)

(3)

将YOLOv5s结构上的特征融合方式结合BiFPN进行修改不仅可以加强特征融合,提高小目标检测的精度,解决航拍图像中的明渠漂浮垃圾大部分尺寸都很小的问题,还可以提高检测速度,确保明渠漂浮垃圾检测的实时性。

2.3 引入注意力机制

为了关注图像中待检测的漂浮垃圾,减少潜在信息的丢失,避免水面倒影、强光反射等对漂浮垃圾检测的干扰,加强漂浮垃圾特征的提取,该文借鉴CBAM的通道注意力机制和空间注意力机制的结构,通过保留跨三维的信息减少信息的丢失,同时使用两个卷积层促进空间信息融合从而使网络可以获得更多信息。CBAM是一个轻量级的通用注意力机制模块,可以直接添加到目标检测网络中进行训练,而且带来的计算开销也极小,不会影响实时目标检测。CBAM注意力机制模块包含两个子模块,一个是通道注意力机制模块,另一个是空间注意力机制模块,其结构如图4所示,图中⊗表示元素级乘法。

图4 CBAM注意力机制结构

虽然CBAM可以使特征图能够更多地包含和表现有效信息,但是它忽略了空间和通道之间的相互作用,会造成跨维信息的丢失。于是,对CBAM的通道注意力机制和空间注意力机制进行修改,称为T-CBAM。首先在通道注意力机制模块中使用3D排列[15]来保留跨三维的信息,从而利用每一对三维通道、空间宽度和空间高度之间的注意力权重来减少信息的丢失,提高检测精度;然后在空间注意力机制模块中使用两个卷积层实现空间信息融合,从而可以关注到更多的空间信息。其具体实现过程如下:

主干网络的特征进行特征融合之后,进入通道注意力机制,先将输入的特征图在通道、空间宽度和空间高度的3个维度上实现全局平均池化和最大池化,并将池化得到的权重矩阵传入权重共享的多层感知器MLP(Multi-Layer Perception)中进行特征相加,最后通过Sigmoid激活函数得到通道注意力特征图Mc(F)。T-CBAM的通道注意力机制如图5所示。

图5 T-CBAM的通道注意力机制

之后,通道注意力机制输出的特征图将作为输入的特征图进入空间注意力机制,同样在通道、空间宽度和空间高度的3个维度上实现全局平均池化和最大池化,并将结果串联连接起来,然后采用两个7×7的卷积核进行卷积操作,保留更多空间信息,最后通过Sigmoid激活函数得到空间注意力特征图Ms(F),实现特征细化。T-CBAM的空间注意力机制如图6所示,其中r是通道注意力机制中MLP使用的压缩比。

图6 T-CBAM的空间注意力机制

特征融合之后引入T-CBAM注意力机制的总体过程可以概括为:

(4)

其中,⊗表示元素级乘法,F'表示经过通道注意力机制输出的特征图,F''表示经过通道注意力机制和空间注意力机制最终输出的精化特征图。

该文经过Neck部分的特征融合之后引入3个T-CBAM,这样会减少潜在信息的丢失,使特征表现更加充分,加强特征提取,抑制无用信息,从而更好地提高检测性能。引入T-CBAM注意力机制可以实现集中对水面漂浮垃圾特征的提取,避免水面倒影及无害漂浮物干扰水面漂浮垃圾检测,而且加强特征提取也会减少水面强光反射对明渠漂浮垃圾检测的影响。

2.4 增加网络的输入尺度

小目标本身像素点少,自身携带的特征信息少,有时很难被发现,导致位置信息获取困难,最直接的解决办法就是增大网络输入分辨率,这样可以使图像具有更加丰富的细节信息和更加精确的位置信息,有利于小目标特征信息的提取。于是,该文将输入尺度从640×640增大到1 280×1 280,增大后的预测分支分辨率则变为32×32,64×64,128×128。这样不仅可以提高目标检测的检测精度,还可以提高小目标的检测精度,也有利于解决航拍图像中的明渠漂浮垃圾大部分尺寸都很小的问题。

3 实验与分析

3.1 实验环境及参数设置

为了训练和评估所提出的模型,实验使用的是CentOS Linux 7.9系统,并采用深度学习框架Pytorch。所有模型都是在GPU上进行训练的,其中GPU型号为Tesla V100-SXM2-16GB。实验环境为Python3.8,Pytorch 1.8.1,CUDA10.1,Cudnn7.6.5。实验使用的数据集是FloW-Img数据集的扩充数据集,即通过剪裁、翻转、加入高斯噪声、改变对比度等方式进行图像处理从而增加数据集的数量。实验采用的是Adam(Adaptive moment estimation)优化器,总共迭代300次,其中批量大小设置为16,初始学习率设为0.01,动量为0.937,权重衰减系数为0.000 5。

3.2 评价指标

为了评价模型的性能,主要采用精确度P(Precision)、召回率R(Recall)、平均精度均值mAP(mean Average Precision)、运行时间(Times)、检测速度FPS(Frames Per Second)、浮点运算量GFLOPs(Giga Floating-point Operations Per Second)和参数量Params(Parameters)作为模型性能的评价指标。

精确度表示正确分类的样本占总样本的比例,如式(5),精确度越大表示检测越精确,在一定程度上也可以反映误检越少。召回率表示正确的正样本占总正样本的比例,如式(6),召回率越大表示检测越全面,反之,说明了漏检越少。mAP表示所有类别的平均精度AP(Average Precision),如式(7),mAP越大表示模型检测精度越大。其中AP是指在不同召回率下精确度的均值,也就是以召回率为横坐标、精确度为纵坐标的P-R曲线与坐标轴围成的面积,如式(8)。mAP_0.5表示IoU阈值为0.5时所有目标类别的平均检测精度。IoU(Intersection over Union,交并比)计算的是预测边框与真实边框的重叠率。运行时间表示模型进行检测消耗的时间,单位为毫秒(ms)。FPS表示模型每秒能检测多少张图像,单位为f·s-1,用于衡量模型的实时性,FPS越大表示检测速度越快,即模型越具有实时性。GFLOPs表示每秒10亿次的浮点运算,用于衡量训练模型时的计算复杂度,GFLOPs越大表示模型越复杂。参数量表示模型训练需要的参数总数,参数量越大所需的内存越大。

(5)

(6)

(7)

(8)

式(5)和式(6)中,TP(True Positives)表示明渠漂浮垃圾被正确检测为漂浮垃圾的个数;FN(False Negatives)表示明渠漂浮垃圾没有被检测出来的个数,即漏检的个数;FP(False Positives)表示将其他的物体检测成漂浮垃圾的个数,即误检的个数。公式(7)中,n表示类别,N表示总类别。公式(8)中,p(r)表示以召回率为横坐标、精确度为纵坐标的P-R曲线。

3.3 对比实验

3.3.1 添加CBAM模块与T-CBAM模块对比

将直接添加CBAM模块与添加T-CBAM模块的检测结果进行对比。它们的检测效果如表1所示,其中YOLOv5s-CBAM表示添加CBAM模块,YOLOv5s-T-CBAM表示添加T-CBAM模块。在复杂度和参数方面虽然不是最佳,但是相差不大,而且从模型的检测精度和检测速度两个方面分析可知,YOLOv5s-T-CBAM的检测效果更好,精确度、召回率和平均精度均值都是最高的,且检测速度仍可以达到89 f·s-1,可以保持较好的实时性。

表1 添加CBAM模块与T-CBAM模块结果对比

3.3.2 与其他目标检测算法的对比

为了验证改进后算法的优越性,还将改进的YOLOv5s算法与YOLOv3、YOLOv4、SSD、YOLOX、DETR、YOLOv7等目标检测算法进行对比。由表2分析可知,在检测精度方面,改进的YOLOv5s算法明显优于其他算法,mAP_0.5可以达到0.899,而且还兼顾高精确度和高召回率,既保证了检测的准确性,还减少了检测的误检和漏检;在运行时间和检测速度方面,改进的YOLOv5s算法的检测速度达到了96 f·s-1,且运行时间最短,满足实时性的要求;在复杂度和参数方面,改进的YOLOv5s算法仅次于YOLOv5s,且相差不大,以最小的牺牲提高了目标检测的精度和速度。因此,改进的YOLOv5s算法优于其他目标检测算法。

表2 改进后算法与其他目标检测算法的对比

3.4 消融实验

为了验证各个改进模块优化的作用,该文进行了消融实验,一共设计了5组实验,其实验结果如表3所示。其中原YOLOv5s表示输入图像的尺寸为640×640,YOLOv5s表示输入图像的尺寸为1 280×1 280。由表3分析可知,在检测精度和速度方面,增加网络的输入尺度提升了检测精度,但是检测速度会有所下降;修改特征融合过程不仅可以大大提升检测精度、精确度和召回率,还可以提高检测速度,弥补增加网络输入尺度时造成的速度损失;引入T-CBAM提升了检测精度和精确度,但是检测速度会降低;在增加网络的输入尺度的同时,既修改网络融合过程,又增加T-CBAM模块,即提出的改进算法,它不仅最大程度地提升了检测精度,使精确度和召回率到达最高,而且检测速度也有了一定提升。在复杂度和参数方面,增加网络的输入尺度不会改变模型的复杂度,参数量也不会增加,但是修改特征融合过程和引入T-CBAM都会影响模型复杂度,增加参数量,通过具体数据对比分析可知,改进YOLOv5s算法的GFLOPs仅比原YOLOv5s算法增加了0.3,且参数量变化不大,以最小的牺牲提升了模型的检测性能。由图7分析可知,改进YOLOv5s的平均精度每次都比原YOLOv5s高,保持实时性的同时在精度方面有很大的提升。

表3 消融实验

图7 原YOLOv5s与改进YOLOv5s的mAP_0.5对比

3.5 检测结果分析

为了验证改进的YOLOv5s算法的可行性,分别使用原YOLOv5s算法和改进的算法对测试集上航拍图像的漂浮垃圾进行检测,以检测瓶子为例,如图8为原YOLOv5s算法和改进的YOLOv5s算法在不同干扰场景下的检测结果对比图,其中,左图为原YOLOv5s算法的检测结果,右图为改进YOLOv5s算法的检测结果。图8(a)表示水面漂浮垃圾较小的检测,可以看到左图会漏检掉远处较小的漂浮垃圾,而右图可以检测到远处的小漂浮垃圾;(b)表示水面倒影干扰的检测,左图将岸边物体在水面的倒影误检成漂浮垃圾,右图则检测正常;(c)表示水面强光干扰的检测,左图受强光干扰,存在漏检,右图不仅将每个漂浮垃圾检测都出来了,而且检测的精度比左图高。综上,改进的YOLOv5s算法更优于原YOLOv5s算法,原YOLOv5s算法在复杂检测场景的干扰下检测性能较差,容易出现漏检和误检的情况,并且相较于改进的YOLOv5s算法,其目标定位精度偏低。因此,改进的YOLOv5s性能更优越,定位精度更准确。

图8 不同干扰场景下算法检测结果对比

4 结束语

针对明渠漂浮垃圾检测存在的一系列问题,如:相关数据集较少,航拍图像中的明渠漂浮垃圾尺寸较小,明渠漂浮垃圾检测易受水面倒影和强光反射影响,基于原始的YOLOv5s进行改进优化。首先,通过数据增强进行数据集扩充,避免数据集太小出现过拟合的问题;然后,在借鉴BiFPN加权双向特征融合的思想修改特征融合过程,提高小目标的检测精度和速度;并且在Neck和Head部分之间添加3个改进的CBAM注意力机制,能够有效降低检测的漏检率和误检率,提高检测精度;最后,增大网络输入的分辨率,使图像具有更加精确的位置信息和更加丰富的细节,有利于小目标特征信息的提取。由实验结果可知,所提方法在保持实时性的同时,检测的平均精度达到了89.9%,具有较好的应用价值。但是,如何将算法部署到嵌入式平台使无人机和无人船能够对视野中的目标进行实时识别是未来亟待解决的问题。