基于深度学习的水面漂浮物识别算法设计

刘 麒，尹 港，王 影，叶 泽

(吉林化工学院 信息与控制工程学院，吉林 吉林 132022 )

近年来，随着我国综合国力的快速发展，湖泊废弃物问题已经引起了研究人员和公众的关注.水面漂浮物如塑料水瓶、牛奶盒、渔网等，容易对水体中的生物造成一定的危害[1].因此，为保护水中生物的安全，需要对水面的垃圾进行清理.如今依然靠人工驾驶大型船只清理水面垃圾，这需要消耗大量的人力和物力.所以，采用水面垃圾清理机器人对水面的垃圾进行检测和清除成为一种有效的方式.

目标识别工具包括合成孔径雷达、红外成像传感器和摄像头等[2-3].随着计算机技术的提高和基于深度学习检测算法的发展，摄像头识别目标物成为一种流行的方式，视觉系统提供的关于物体外观的数据可以得到更好地利用.为了水面垃圾清理机器人能够实时且准确识别，用摄像头配合深度学习检测目标物是高效的解决方案.

1 算法原理

YOLO算法是Joseph Redmon等人在2016年发表的[4-5]，随着计算机技术的进步和深度学习框架的发展，YOLO算法已经发展到第五代.YOLOv5目标检测算法是Ultralytics公司于2020年发布的，YOLOv5有4个版本，包括YOLOv5s,YOLOv5m，YOLOv5l，YOLOv5x.由于要将算法配置在树莓派上，要求检测速度要快，所以选择了模型较小的YOLOv5s.图1是YOLOv5目标识别算法的结构图.

YOLOv5基于CSP-Darknet网络，输入一张608*608三维度的图像，经过主干网络对图像进行卷积和切片等操作，将主干网络的信息传到Neck网络.Neck网络将这些图像信息进行融合、上采样和下采样等，然后输出3种尺度的图像，76*76*255尺度的图像用来预测小目标；38*38*255尺度的图像用来预测中等目标；19*19*255尺度的图像用来预测大目标[6].

输入端主要对输入的图片进行修改，进行模糊处理，几张图片进行缝合，或对图片进行随机排列，使得输入的图片统一缩放到某个标准尺度，然后再送入检测网络.主干网络主要包括CSP-Darknet网络，有减少网络计算量以及对显存的占用，同时保证网络的能力有所提升的优点[7].Neck模块主要包括改进的SPP结构和CSP-PAN结构，其中改进的SPP结构是将输入串行通过多个5*5大小的最大池化层，进一步增大特征图的感受野.CSP-PAN结构使得物体在不同尺度下输入时可以检测得较好.输出端对检测框和目标框进行损失函数的计算，采用非极大值抑制用于预测框的筛选，然后输出带有预测框的图片.

2 改进思路

2.1 YOLOv5算法中K-means算法改进

YOLOv5算法采用K-means聚类算法生成anchor boxes.原始YOLOv5算法中先验框的尺寸是由K-means算法在COCO数据集上聚类得到.K-means聚类算法的核心思想是：数据集中随机生成K个聚类中心点，计算数据集中其余点到聚类中心点的欧式距离，与聚类中心点距离较小的点划分到一组，从而完成对数据集的聚类.但由于K-means算法对生成聚类组采用的欧氏距离计算结果并不是太好，所以本文采用的是数据集中随机生成的聚类中心点的先验框与其余先验框的重合程度来表示，即IOU越大，表示距离越接近.从图2可以看出，IOU是一个比值，即交并比.在分子中计算数据集聚类中心先验框与数据集中其余先验框的重叠区域，分母是两个先验框的并集区域.改进的K-means算法聚类由公式(1)、(2)所示.使用改进的K-means聚类算法对数据集中的数据进行聚类，并经过200次遗传算法迭代，生成3组anchor boxes：[5,4,9,5,10,7]，[15,6,15,10,24,9]，[29,12,39,18,60,34]，使得模型更容易收敛，提高检测的精度.预测框与真实框的IOU表示如图2所示.

(1)

D=1-IOU.

(2)

2.2 SE注意力机制模块

SE注意力机制在通道维度增加注意力机制,原理图如图3所示.通过自动学习的方式，使用新的神经网络，获取到特征图的每个通道的重要程度，然后用这个重要程度去给每个特征赋予一个权重值，从而让神经网络重点关注某些特征通道.提升对当前任务有用的特征图通道，并抑制对当前任务用处不大的特征通道，这样增加了参数和计算量，也增大了有效信息的权重，使得网络更好训练[8].

在输入SE注意力机制之前，特征图每个通道的重要程度都是一样的，通过SE网络之后，不同颜色代表不同的权重，使每个特征通道的重要性变得不一样了，神经网络重点关注某些权重值大的通道.Squeeze通过全局平均池化，将每个通道的二维特征(H*W)压缩为一个实数，将特征图从[h,w,c]转换为[1,1,c],如公式(3)所示.

(3)

Excitation给每个特征通道生成一个权重值，如公式(4)所示.

s=σ(W2δ(W1z))，

(4)

Scale将归一化权重加权到每个通道的特征上，如公式(5)所示.

(5)

2.3 损失函数改进

边界框预测是目标检测中重要的任务之一，在图像中想要得到检测出的目标，就要对边界框所在的位置进行预测.边界框预测损失函数包括IOU_loss、GIOU_loss、CIOU_loss等.其中IOU_loss主要考虑真实框和预测框的重叠面积，但是当真实框和预测框没有相交时，这就造成边界框损失函数计算不准确.在IOU_loss的基础上，GIOU_loss考虑了边界框不相交的问题.CIOU_loss则考虑了预测框和真实框中心点之间的距离，用欧氏距离计算两点之间的距离，另外还增加了对边界框宽高比的考虑,如公式(6)所示.α-IOU通过调节参数α，使预测框有很大灵活性，对小数据集和噪声有更强的鲁棒性，如公式(7)所示.

YOLOV5目标检测算法中边界框预测采用GIOU_loss，本文采用α-CIOU作为边界框预测目标的损失函数.如公式(8)所示.

(6)

Lα-IOU=1-IOUα1+Ρα2(B,Bgt)，

(7)

(8)

b代表了预测框的中心点；bgt表示真实框的中心点；ρ(·)表示两个中心点之间的欧式距离；c表示预测框和真实框的对角线距离；α表示可调节参数.

3 模型训练

3.1 程序运行流程

基于改进YOLOv5深度学习的水面垃圾识别算法主要包括建立水面垃圾数据集、水面垃圾图像标注和网络训练等环节，具体流程如图4所示.

3.2 实验平台

实验操作系统为Ubuntu，CPU型号为Intel(R) Xeon(R) Gold 5320 CPU @ 2.20GHz，运行内存为32GB.GPU型号为RTX A4000，显存大小为16GB.该实验在Pytorch 1.7.0深度学习框架、Cuda 11.0架构上实现，Python 3.8为基础编程语言.训练的参数：batch_size设置为32，训练周期设置为200，学习率设置为0.001，其他参数为默认值.

3.3 实验数据集

本文使用的数据集是欧卡智能提供的公开数据集(FLOW)[9]，该数据集适应于内陆水域复杂场景，且是从水面无人航行器的角度收集的水面漂浮物.其他数据集，如包含野外废弃物的TACO数据集[10]，用无人驾驶飞行器(UAV)拍摄的地面上真实废弃瓶的UAV-DB数据集[11]，还有在室内拍摄的废弃物MJU-Waste数据集[12].本文使用的数据集中超过一半的标记漂浮物被视为小物体，与其他数据集相比，更适合对水面漂浮物识别进行训练，可以促进对水面上小物体检测的研究.如图5所示，(a)、(b)、(c)分别是TACO数据集、UAV-DB数据集和MJU-Waste数据集，(d)是本所使用的数据集FLOW，尽管(a)、(b)、(c)数据集对废物分类很有价值，但是这些数据集应用场景与内陆水域有很大的不同.

由于原始数据集只有500张图片，数据集的数量也是决定目标检测效果的一个重要因素，为了得到更好的目标检测效果，需要对数据集进行扩充.扩充方式包括对原始图像加入高斯模糊和旋转180°等.经过两种扩充方式之后，数据集由原有的500张图像增加到1 500张图像，如表1所示.原始数据集与扩充之后的数据集对比如图6所示.

表1 数据集扩充数量

3.4 评价指标

本实验引入查准率(Precision)、查全率(Recall)、平均精度AP(Average Precision)和平均精度均值mAP(mean Average Precision)来评估检测模型的准确性.其中TP代表检测到水面漂浮物，并且在实际图片中也检测到了水面漂浮物；FP代表检测到水面漂浮物，但是在实际图片中未检测到水面漂浮物；FN代表未检测到水面漂浮物，但是在实际图片中检测到水面漂浮物.

查准率是模型预测的所有目标中，预测正确的比例，如公式(6)所示：

(6)

查全率是所有真实目标中，模型预测正确的目标比例，如公式(7)所示：

(7)

AP(Average Precision)表示P-R曲线下的面积；mAP(mean Average Precision)表示各类别AP的平均值.因本文使用的数据集中的类别只有bottle一项，所以mAP值与AP值相等.

4 实验结果

原始YOLOv5模型和改进的YOLOv5模型在训练过程中，均使用相同的参数设置和数据集，根据训练过程中的日志文件绘制出如图7所示的损失函数曲线.将训练集得到的权重文件应用到验证集上进行验证，通过表2可以知道改进的算法比原始算法在查准率和平均精度均值分别提升了2%和4%.将改进的YOLOv5算法与原始YOLOv5算法在除本文所使用的数据集之外的图片上进行对比，改进的YOLOv5算法在误检率、检测精度方面优于原始YOLOv5算法，如图8所示.

表2 实验结果

5 结 论

引入改进的YOLOv5目标识别算法，在水面漂浮物尤其是小目标检测方面有了较好的结果.首先，通过改进的K-means算法对数据集的先验框重新聚类，提高了模型的收敛性.其次，将α-IOU损失函数与c-IOU损失函数结合，为小目标检测提供了更强的鲁棒性.改进后的YOLOv5算法与原始算法相比，在查准率和平均精度均值方面分别提升了2%和4%.