基于小目标检测的YOLOv5的静止水面垃圾检测与分类算法

Detection and Classification Algorithm of Static Water Garbage Based on YOLOv5 Small Target Detection

摘要: 针对现有水面垃圾检测模型精确率低、运算速度慢以及鲁棒性差的问题，提出一种基于改进型YOLOv5的水面垃圾检测与分类算法，将小目标检测头引入改进型网络框架以提高目标检测准确率，通过融入DeepSORT算法以提升检测速度，加入SENet注意力机制并更改损失函数为SIoU以提高模型的鲁棒性。通过对洞庭湖水域垃圾检测结果展开改进前后对比分析，实验结果表明，与原始YOLOv5相比，改进后的YOLOv5的mAP@0.5达到92.1%，较改进前提高了5.37%，同时平均帧率提高至47.8FPS，较改进前提升19.20%，为后续开展水面垃圾自适应处理可行性方案奠定了坚实基础。

Abstract: In response to the problems of low accuracy, slow operation speed, and poor robustness of existing water surface garbage detection models, a water surface garbage detection and classification algorithm based on improved YOLOv5 is proposed. The improved network framework introduces a small target detection head to improve target detection accuracy, incorporates the DeepSORT algorithm to enhance detection speed, adds SENet attention mechanism and changes the loss function to SIoU to improve the model's robustness. Through comparative analysis of the garbage detection results before and after improvement in the Dongting Lake area, the experimental results show that the improved YOLOv5 achieves an mAP@0.5 of 92.1%, an increase of 5.37% compared to before improvement, and the average frame rate is increased to 47.8 FPS, an increase of 19.20% compared to before improvement, laying a solid foundation for subsequent exploration of the feasibility of adaptive processing of water surface garbage.

1. 引言

随着城乡现代化的脚步大步向前迈进，旅游业已经成为中国第三大国民经济支柱，而其中的水面湖泊旅游风景区更是数量极多，尤其是在我国东部地区，存在着大量地势差较小的水面湖泊风景区，因没有得到很好的管理 [1] ，而在出水口等处，堆积大量垃圾，导致本湖泊风景区水质下降，并且对后续下游水域产生二次污染，这种现象对水质和水生态系统造成了很大威胁，同时也严重影响风景区的观赏性 [2] ，从而减少了旅游带来的经济效益 [3] ，从更深层次来看，是对国民经济 [4] 造成了冲击。国家相关部门强调“生态环境保护和经济发展是辩证统一、相辅相成的 [5] ，建设生态文明、推动绿色低碳循环发展，不仅可以满足人民日益增长的优美生态环境需要，而且可以推动实现更高质量、更有效率、更加公平、更可持续、更为安全的发展，走出一条生产发展、生活富裕、生态良好的文明发展道路”。“生态环境也是关系民生的重大社会问题”，所以保护水资源也十分必要，亟不可待。目前，虽然政府也正大力推动水资源保护，许多的水域保护政策也相应出台，例如“河长制” [6] ，但实际上水面垃圾的监测和清理方式主要还是依靠人工巡查或使用大面积筛网 [7] ，这些方式不仅效率低下，且受人员限制，只能覆盖到小范围水域 [8] ，虽然市面上偶尔会开发出新兴的水面清理机器人 [9] ，但其清理方式因为缺少更远见的指引和更合理的清理决策，大多也只是漫无目的、大小通吃的全部清理，导致能源利用效率低，浪费大量能源。据此来看，水域垃圾清理问题已成为环境保护和经济回暖的一个亟待解决的问题。

本团队以洞庭湖南湖水域水情为基点，开展了具体研究，发现了以下现象及问题。游客产生的垃圾不断堆积，最终漂浮在水面，形成静止水面垃圾，只依赖人工的清理方式会加剧财政压力，并且清理效率低下，所以现在主要采取的是结合YOLOv5检测系统辅助清理垃圾，传统的YOLOv5并未引入小目标检测头。因此我们针对传统YOLOv5检测算法进行改进。

传统的水面垃圾检测算法对于不同的场景变化检测准确率变化较大，不能满足检测水面垃圾的精度要求。为了提高检测精度早期杜艾卿 [10] 等使用GhostBottleneck模块替换YOLOv5s模型结构中的C3和部分CBS模块，并加入ECA注意力机制，但最高的精准率为90.3%，并且目前水面垃圾检测对小目标的准确识别精度低，还不能满足对静水水面垃圾的检测要求。王一早 [11] 等对YOLOv3的K-means算法进行了改进，使M_IoU提高了12.8%，M_IoU是目标框的最大IoU值。也有其他相关研究 [12] 也采取了数据增强的方式去提高检测精度，但大多效果一般。

本文以动态视频中的静止水面垃圾为检测对象，对传统的YOLOv5算法进行改进，除了对注意力机制和损失函数 [13] 方面的基础改进，也注重了对小目标检测 [14] 的关注，本文把改进后的YOLOv5称为YOLOv5s-self。通过应用先进的算法，计算机视觉技术 [15] ，实现快速、准确定位与检测静止水面垃圾，提高了对静止水面的检测精度。我们将YOLOv5融合了DeepSORT算法，可以对动态视频中的水面垃圾进行更好的追踪，同时我们添加了SENet注意力机制，来提高模型的鲁棒性和泛化能力 [16] ，通过将损失函数GIoU替换成SIoU，提高了预测框的精度，添加小目标检测头，能够更好的提取小目标的特征信息，从而提高对小目标的检测的准确率。

2. 改进YOLOv5算法

为了准确检测动态视频中的静止水面垃圾，我们基于YOLOv5进行改进，拟开展以下对YOLOv5程序的改良技术路线的研究，如图1。

Figure 1. Improved technical route of YOLOv5

图1. YOLOv5的改良技术路线

Figure 2. Improved network framework of YOLOv5s-self model

图2. 改进YOLOv5s-self模型的网络框架

本文先将YOLOv5融合DeepSORT算法 [17] ，可以增加对目标的追踪，可以有效的增加检测速率和精确度。再加入小目标检测头 [18] ，让识别标签能紧跟识别物，增强对小对象的特征提取能力，最后为了进一步提高精度，我们改变网络结构 [19] ，将损失函数更改为SIoU [20] 并添加SENet注意力机制 [21] ，增强对于对象检测的准确性以及检测速度，对于湖泊垃圾检测与分类精准度和速度都有显著提升，具体改进过程如图2。

2.1. 添加DeepSORT算法

针对动态视频中湖泊垃圾检测和分类可能出现的检测速度慢和精度差问题，融合DeepSORT算法可以有效的应对这类问题。DeepSORT是一种经典的多目标跟踪算法，它是一个两阶段的算法，可以实现实时跟踪效果。DeepSORT基于Sort目标跟踪进行改进，引入了深度学习模型，在实时目标跟踪过程中，提取目标的外观特征进行最近邻近匹配。具体来说，DeepSORT首先使用卷积神经网络(CNN)提取每个目标的外观特征，然后使用匈牙利算法将当前帧中的每个检测框与上一帧中的跟踪目标进行匹配，得到匹配分数。接着，使用卡尔曼滤波器对跟踪目标的位置和速度进行预测，并根据匹配分数和预测结果对跟踪目标进行更新和筛选 [22] 。这可以对小目标进行更加准确的跟踪 [23] ，避免了传统算法中容易出现的漏检或误检问题，并可以在较短的时间内完成目标跟踪，提高了跟踪的效率。DeepSORT流程如下图3。

Figure 3. DeepSORT algorithm framework

图3. DeepSORT算法框架

2.2. 添加损失函数SIoU

原始的YOLOv5使用的是GIoU (Generalized IoU)作为目标检测任务的损失函数。SIoU在计算损失时对目标框的重叠程度进行平滑。它是在GIoU损失函数的基础上进行改进得到的。平滑参数可以控制损失函数对重叠程度的敏感度，从而使模型对于不同程度的重叠情况具有更好的鲁棒性。SIoU损失函数考虑了目标框的结构信息，可以更好地评估目标框的匹配程度。SIoU损失函数的计算公式如下：

$\text{SIoU=IoU}-\text{ln}\left(\text{1+}\alpha \cdot \frac{V}{S}\right)$ (1)

其中，IoU，V，S，α分别是目标框的IoU值，目标框的体积，目标框的表面积，一个超参数，范围是0到1，用于平衡体积和表面积的影响。SIoU考虑到期望回归之间向量的角度，重新定义角度惩罚度量，它可以使预测框快速漂移到最近的轴，随后则只需要回归一个坐标(X或Y)，这有效地减少了自由度的总数。

2.3. 添加SENet注意力机制

为了帮助模型更好的处理复杂场景下的大规模数据，我们在backbone层增添了SENet注意力机制，SENet结构图如下图4。

Figure 4. SENet structure

图4. SENet结构

F_sq，H，W，U分别是全局均值池化，输入数据的高，输入数据的宽度，输入数据。

F_ex，W₁，W₂，σ分别是两级全连接，学习得到的权重矩阵，sigmoid激活函数。

在YOLOv5中，通常使用的通道数是3 (RGB图像的三个通道)，但是当引入SENet时，在每个模块之后添加一个SE块，其中包含一个全局平均池化层和两个全连接层。这些层用来学习通道之间的相关性并获得通道的权重信息。它通过挤压和激励操作来自适应地调整每个通道的重要性，从而提高网络的性能 [24] 。具体来说，SENet引入了一种称为“Squeeze-and-Excitation block”的模块，该模块包括两个步骤：

(1) Squeeze：通过全局平均池化操作将每个通道的特征图压缩为一个数值，得到一个通道描述符向量。

(2)

(2) Excitation：通过两个全连接层来学习每个通道的权重，从而自适应地调整每个通道的重要性。

$s=F_{ex}\left(z,W\right)=\sigma \left(g\left(z,W\right)\right)=\sigma \left(W_2\delta \left(W_{1}z\right)\right)$ (3)

SENet可以自适应地学习每个通道的重要性，从而提高模型的准确性，并可以通过学习每个通道的权重来减少模型的参数，从而提高模型的效率，同时可以应用于各种计算机视觉任务，例如分类、检测、分割和生成模型等。

2.4. 添加小目标检测

针对动态视频中的湖泊垃圾的小对象像素点较少导致的漏检问题，我们添加小目标检测头来提高对图像中小对象的特征提取能力。通过在网络中引入多个尺度的特征图，并进行融合操作，可以提供更全面、多样化的特征表达能力，以适应小目标的多样性和复杂性 [25] 。由于小目标的像素数通常较少，为了更好地捕捉和识别小目标，增加卷积层的数量或者使用特殊的卷积核来保留更细节的信息用于提取高分辨率特征的操作 [26] 。同时引入上下文模块，以增强模型对小目标的感知能力，这有助于提升小目标边界的清晰度和准确度，检测模块结构如下图5。

Figure 5. SSH detection module

图5. SSH检测模块

通过1 × 1的卷积核可以将多个通道的特征图进行线性组合，融合不同层级的信息，增强模型对特征的表达能力，上下文模块如下图6。

Figure 6. SSH context module

图6. SSH上下文模块

上下文模块通过引入额外的上下文信息，扩展了神经网络的感知范围。传统的卷积神经网络只考虑了局部信息，而上下文模块能利用前一层的特征信息，并将其传递到后续层中，并且能够捕捉更广阔的上下文关系 [26] 。

2.5. 模型算法评价指标

目标识别检测算法，判断所识别目标的识别精度的主要指标之一是Map (英文全称mean average precision)、Precision (精准率)、召回率(Recall)，计算力(GFLOPS) [27] 。如下为计算公式：

$Precision=\frac{TP}{TP+FP}\times 100\text{\%}$ (4)

$Recall=\frac{TP}{TP+FN}\times 100\text{\%}$ (5)

$mAP=\frac{{\displaystyle {\sum }_{i=1}^{n}A{P}_i}}{m}$ (6)

$GFLOPS=\left(浮点运算次数\right)/\left(以秒为单位的运行时间\right)/\left({10}^9\right)$ (7)

其中TP，FP，FN分别为模型将负样本错误检测为正样本的个数，模型正确预测正样本的个数，正样本被错误识别为负样本的个数。Recall为召回率，即模型正确检测出正样本的个数所占的比率，m为类别的数量， $A{P}_i$ 是精度与召回率的曲线积分。

3. 性能评估

3.1. 实验环境

硬件平台CPU：AMD Ryzen 7 5800H with Radeon Graphics 3.20 GHZ；内存：16.0 GB；存储空间：512 GB。

软件平台：开发环境PyCharm 2020.1.1 (Professional Edition，编程语言Python 3.9)。

3.2. 结果分析

我们对3种垃圾进行了多次测试，训练集和测试集的比例选定为10:1，训练了我们的数据集，其中的数据分布具体见表1。

Table 1. Distribution and quantity of data sets

表1. 数据集分布与数量

通过我们对YOLOv5算法的改进，我们测试的结果也有了很大的提升，改进前后精度–置信度由0.764提升至0.859，提升了0.095，召回率也有显著提升，可以高效且高质量完成水面垃圾分类任务，精准的将水面垃圾分为了3类，分别是瓶罐类，小垃圾类和死亡鱼类。

改进前mAP@0.5为0.874，改进后提升至0.921，提升了0.047，具体见图7。

Figure 7. Comparison of detection values before and after

图7. 检测值前后对比

分析上述实验结果，可以看出YOLOv5-self与YOLOv5s相比，模型参数量增加了9.40%，但是计算力增加了69.62%，增加了CPU利用率，加快了运行速度，减少了运行时间，降低了时延，提升了FPS，精度提高了12.43%，这是因为增加了小目标检测头，提高了计算力，对小对象的特征的提取更充分，可以更精确地识别小目标，具体提升见下表2。

W——Weight file size (权重文件大小)，Par——Parameters (参数量)，GFLOPS (计算力)，mAP@0.5，P——Precision (精准率)，FPS (每秒帧率)。

Table 2. Comparison of models before and after improvement

表2. 改进前后的模型对比

3.3. 消融实验

最后为了验证改进算法的有效性，本文做了多组实验，在使用相同数据集的时候，通过控制变量的实验方法，实验结果对比如下表3所示。

Table 3. Model improvement

表3. 模型改进

由表3可知，在增加小目标检测头，DeepSORT算法，SENet并且用SIoU替换IoU损失函数，改进后，精准率P，召回率R，平均精度均值mAP@0.5" target="_self">mAP@0.5与主流YOLOv5模型相比，均得到不同程度的提升，精准率提高了12.43%，mAP@0.5" target="_self">mAP@0.5提升了5.37%，计算力增加了69.62%。

Figure 8. Comparison between detection chart and atomization treatment results

图8. 检测图与雾化处理结果对比

我们通过对动态视频进行噪音雾化处理后再次进行检测，发现对结果的影响不大，说明本文对于YOLOv5的改进可以很好提高鲁棒性和泛化能力，具体见图8。

4. 总结

本研究在动态视频中的湖泊垃圾检测与分类领域中引入了改进的YOLOv5算法，结合DeepSORT算法进行优化，以解决多个挑战。针对传统图像检测方法的低精度、长处理延迟、有限的垃圾分类准确性以及对光照和天气条件的敏感性等问题，本算法采用了多种改进措施，如整体损失函数优化、引入小目标检测头和注意力机制等，精确度提高了12.43%，mAP@0.5提高了5.37%，模型的计算力提高了69.62%。这些改进显著提高了识别准确性，同时降低了人力和时间成本，实现了大规模成本节约。此外，本研究提供了宝贵的多源图像融合见解，提升了整体系统性能和可靠性。使用改良后的YOLOv5算法进行湖泊垃圾识别和分类，实现了对湖泊垃圾的实时动态监测和识别，为垃圾清理决策提供了参考，可以高效且高质量完成水面垃圾分类任务，精准的将水面垃圾分成了瓶罐类，小垃圾类、死亡鱼类等3类，对水面垃圾的检测达到了较高精度，为后续开展水面垃圾自适应处理可行性方案奠定了坚实基础。

致谢

首先，我要感谢我们的指导老师，您在整个研究过程中给予了我宝贵的指导和悉心的培养。感谢您对我们学术上的启发和教诲，使我们不断成长和进步。

同时，我们要感谢实验室的其他同学，在日常讨论和交流中给予了我们很多灵感和启发，您们的意见和建议推动了我们的研究工作。

最后，我要感谢我们的家人和朋友。感谢他们一直以来对我们的支持和理解，使我们能够专注于研究工作。没有他们的鼓励和陪伴，我们将无法完成这个研究。

在整个研究过程中，有很多人和机构给予了我们无私的帮助和支持，在这里一并向他们表示衷心的感谢。

基金项目

2023年度湖南省大学生创新训练计划项目：基于无人机航拍和改进的YOLOv5的水面垃圾检测与分类系统(S202310543043)。

附录

部分数据代码展示：

损失函数：

bias = 0.5# bias

Offsetssss = torch.tensor([[0, 0],

# [1, 1], [1, -1], [-1, 1], [-1, -1],#jk,jm,lk,lm

], device=targets.device).float() * g# offsets

# Offsets

grid = t[:, 2:4]# grid xy

inverse = gain[[2, 3]] - grid# inverse

a, b = ((grid % 1. < g) & (grid > 1.)).T

c, m = ((inverse % 1. < g) & (inverse > 1.)).T

a = torch.stacb((torch.ones_libe(a), a, b, c, m))

t = t.repeat((5, 1, 1))[a]

off = (torch.zeros_libe(grid)[None] + Offsetssss[:, None])[a]

grid = t[:, 2:4]# grid xy

gij = (grid - off).long()