基于改进YOLOv5s的海下垃圾检测算法

吴观茂,王 涛

(安徽理工大学 计算机科学与工程学院,安徽 淮南 232001)

随着工业化快速发展,海洋生态环境遭受了一定程度的破坏,海洋垃圾逐渐引起了人们的关注。海洋垃圾不仅是指水上各类漂浮物,还有不易人类察觉的水下垃圾,给鱼类、植物、藻类等造成了危害。由于海洋面积辽阔,传统海下垃圾检测方式需要付出极大的人力、物力以及时间,且效率较低。

深度学习作为机器学习的重要组成部分,被广泛地应用在目标检测领域[1-2]。基于深度学习的目标检测算法可以分两类[3]。一类是以R-CNN、Fast R-CNN、Faster R-CNN、SPP-Net等为主的二阶段算法,主要原理是先在原始图片上选择感兴趣的区域,即候选区域,然后在候选区域进行目标分类和边界框回归,检测精度比较高,但计算成本大、速度比较慢。另一类是以YOLO、SSD为主的一阶段算法。不同于二阶段算法,其采用端对端的方式直接进行分类和回归,可以实现实时检测,但检测精度比较低。

海洋垃圾检测与许多视觉物体检测问题一样,环境中的微小变化可能会导致物体外观的巨大变化,光的变化以及水的浊度变化都可能会使物体难以或完全无法被检测到。因此,本文提出一种基于改进YOLOv5s的实时目标检测算法,通过分析YOLOv5s算法和TrashCan海下垃圾数据集,对网络模型和损失函数进行优化,以期提高算法对目标的检测性能。

1 YOLOv5算法

YOLOv5算法使用参数depth_multiple和width_multiple来调整骨干网络的宽度和深度,得到YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x 4个版本的模型。其中,YOLOv5s是最简单的模型,检测速度比较快,但精度要低于其他3个。YOLOv5s的网络结构如图1所示。

图1 YOLOv5s的网络结构

假定输入图片大小为640×640。在主干网Backbone部分,YOLOv5首先通过Focus模块对输入的图片进行切片处理以及卷积核为1×1、步长为1的卷积处理,使输入的图片变为320×320,然后利用CSPDarknet53模块进行特征提取。当提取到最深层特征时(如20×20),Neck部分将此时的特征图进行上采样,把特征图的长与宽拉长2倍(变为40×40),与主干网络的40×40特征图进行融合,重复此类操作,可以使自顶向下FPN中3个特征图都包含浅层特征和深层特征。由于提取的特征图经过FPA模块传达强语义特征,不利于目标定位操作,因此加入自底向上的PANet模块[4],使特征图强定位特征,两者强强结合,提高检测精度。Head部分将检测层分为80×80、40×40、20×20 3个不同尺度,在小、中、大目标检测上可以取得很好的效果[5]。

2 改进的YOLOv5s算法

由于YOLOv5s初版算法模型的网络结构与损失函数专门为通用场景而设计,在一些特定场景上应用效果较差。但其网络结构简单清晰、检测精度比较低,有很大改进空间。因此,本文对YOLOv5s提出3处改进,以解决在TrashCan海下垃圾数据集上的检测问题。改进后YOLOv5s的网络结构如图2所示。

图2 改进后YOLOv5s的网络结构

2.1 添加注意力模块

注意力机制使神经网络能够准确地关注输入的所有相关元素,是提高深度神经网络性能的重要组成部分。Hu等[6]提出的挤压和激励(SE)模块是一种极轻的即插即用模块,嵌入ResNet50主干后,其在ImageNet数据集的检测精度提高了1.51%,额外计算增加了约1%。Zhang等[7]将SE模块添加到YOLOv5s中,在VisDrone-2020数据集上的mAP由13.97%提升到14.03%。Wang等[8]提出ECA通道注意力机制,实验结果表明ECA模块效率更高,性能优于SE模块。因此,鉴于嵌入ECA模块带来的性能大于增加的计算成本,本文将ECA模块嵌入CSPDarknet53主干中。ECA模块结构如图3所示。

图3 ECA模块结构

ECA模块在不降低维度的情况下进行逐通道全局平均池化之后,通过考虑每个通道及其k个近邻来捕获本地跨通道交互[9]。输入的特征图xij的大小为HiWj,经过全局平均池化(GAP)后,特征向量的输出为1×1×C,即g(x),计算公式为：

(1)

式(1)中,x为特征值的输入;H为特征图的高;W为特征图的宽。

GAP操作后,进行卷积核为自适应k大小的1维卷积操作,并经过Sigmoid激活函数得到各个通道的权重w,计算公式为：

(2)

w=σ(C1Dkg(x))

(3)

式(2)～(3)中,C为特征图的通道数;C1D为一维卷积;odd代表取奇数;σ代表Sigmoid激活函数。

最后,将权重与原始输入特征图对应元素相乘,得到最终输出特征图。

2.2 构造四级多尺度特征融合

卷积神经网络越深,特征图的感受域越大,每个神经元包含更多的全局和高层语义特征,但会丢失局部特征和细节特征[10]。相反,当卷积神经网络较浅时,特征图神经元中包含的特征往往更局部和详细[11]。由于YOLOv5s的先验框自适应,当输出的特征图包含更多信息时,先验框可以更好地进行分类与回归,所以增加更低一级多尺度融合层,有利于提高小目标的检测。因此,本文采用四级多尺度特征融合层(Four MFF)进行特征值提取,并采用FPN+PAnet组合作为Neck的模型结构。四级多尺度特征融合结构如图4所示。

图4 四级多尺度特征融合结构

在原初YOLOv5中,主干网3次下采样得到80×80、40×40、20×20的特征图与FPN 2次上采样得到的特征图进行多尺度融合,但没有完全利用主干网底层特征图。这是因为YOLOv5为减少计算成本,提高检测速率,只利用了主干网后3次下采样,因此原初YOLOv5对大、中目标识别效果较好,小目标识别效果较差。在三级多尺度融合层的基础上,增加更低一级多尺度融合层,即将FPN、PAnet分别再进行1次上采样,接着将主干网的4次下采样得到160×160、80×80、40×40、20×20的特征图与FPN 3次上采样得到特征图进行融合,使FPN中4个特征图都包含浅层特征和深层特征,然后将FPN中4个特征图与PAnet 3次上采样得到特征图进行融合。这样就可以在融合的特征图上检测,以确保每个对象的预测能充分利用所有特征层的信息。

2.3 引用EIoU损失函数

边框回归是目标检测的关键步骤,选用EIoU作为YOLOv5s算法的损失函数,EIoU的长宽比例损失αv为：

(4)

式(4)中,Aw和Ah为预测框A的宽度和高度;Bw和Bh为预测框B的宽度和高度;Cw和Ch是覆盖2个最小外接框的宽度和高度。

为更有利于模型的优化,引入Focal Loss优化边界框回归任务中的样本不平衡问题,计算公式为：

(5)

式(5)中,a和b分别为A、B的中心点;c为C覆盖的对角线长度。

从式(5)可以看出,EIoU将损失函数分成了IoU损失、距离损失、边长损失,是针对CIoU损失函数改进得来的,避免了当预测框的长与宽满足[(w=kwgt,h=khgt)|k∈R+]时,CIoU中添加的相对比例的惩罚项不起作用。此外,还防止了当w和h其中一个值增大时,另外一个值必须减小,不能保持同增同减的问题。

在目标检测领域中,模型收敛的速率与高低质量样本有关,这是因为大部分样本的真实框与根据锚点得到预测框的IoU不大,此类样本称为低质量样本,且在低质量样本上进行训练时,会导致损失值更容易剧烈波动。因此,EIoU引入Focal Loss可以很好地解决高低质量样本类别不平衡的问题。

3 实验与分析

3.1 目标检测数据集

TrashCan数据集由带注释的图像(目前有 7 212 张)组成,包含对垃圾、人造物体(ROV)和各种海底动植物三大类的观察,共有22个类用于标记对象。图像来自日本海洋地球科学技术局策划的J-EDI数据集。

3.2 配置环境与评估指标

改进的YOLOv5s模型运行环境为阿里AutoDL算力市场的RTX 3090 GPU,配置环境为1.10版PyTorch、3.8版Python、11.3版Cuda。采用改进的YOLOv5s网络结构进行训练,初始学习率设置为0.001,批量大小为8,学习次数为300。

改进后的模型通过查准率(Precision)、查全率(Recall)、 平均精度均值(mAP)3个指标进行评估,其中mAP的计算公式为：

(6)

式(6)中,i为当前类别;AP(i)为当前类别的平均精度;n为类别总数。

3.3 消融实验

为了验证3种改进方法的有效性,以YOLOv5s模型为对照组,设置4个模型,M1：ECA;M2：Four MFF;M3：EIoU;M4：ECA+ Four MFF +EIoU,进行消融实验测试。实验保留需要测试的改进方法,并屏蔽了可能影响测试模块的其他方法。消融实验结果见表1。训练结果指标变化如图5所示。

表1 消融实验结果

(a) Precision (b) Recall (c) mAP50

从图5和表1可以看出,每种改进方法都比原始YOLOv5s模型的效果好,模型的mAP从原来的87.3%分别增加到88.6%、90.5%、89.4%、92.1%,说明从3个方面改进可以使模型专注于感兴趣的特征区域,并抑制不感兴趣的区域,保留了更丰富的对象特征信息,并帮助损失函数更好地收敛。

改进前后YOLOv5s算法可视化的对比图如图6所示。从图6可以看出,改进后的YOLOv5s无论从目标对象识别程度还是回归框的准确度,都要优于改进之前的YOLOv5s,验证了本文改进算法的有效性。

(a) 改进前 (b) 改进后

4 结论

基于YOLOv5s提出了一种应用于海下垃圾检测任务的改进算法,在主干网中添加ECA模块使模型能够关注感兴趣的区域。为了使特征图的信息更加丰富,在Neck部分新添加一层多尺度特征融合层,引用EIoU作为边界框损失函数,不仅可以解决目标框与预测框的不相交、包含问题,而且能反馈更加合适和真实的预测框位置。本算法在保持原有三层检测层不变的基础上,有效提高了YOLOv5s在TrashCan数据集检测中的准确率、召回率和平均精确度,满足任务要求,解决了融合过程中的特征信息丢失造成的一些目标类别精度下降难题。