基于YOLO v5与层次化分类算法的生活垃圾识别研究

吴 超 蒋鹏飞 吕 刚

（1.合肥学院生物食品与环境学院；2.合肥学院先进制造工程学院 安徽合肥 230601）

随着生活垃圾分类立法、执法等工作的顺利开展以及人工智能的飞速发展，视觉处理中的垃圾目标检测对垃圾处理工厂提升分拣效率、降低工人劳动强度等方面有着重要的研究意义。近年来基于深度学习中目标检测准确率高、模型学习速度快等特点，⓪其在日常生活中的应用场景越来越广泛，众多学者也转向了目标检测的研究领域。深度学习目标检测根据设计理念的不同可分为有无anchor[1]两大类：在anchor出现前，最初R-CNN[2]算法使用的Selective Search[3]方法以及滑动窗口等方法，其中Selective Search方法是能够分割出多达上千个候选区域，依据图像颜色、纹理、合并后的最小总面积等原则生成两个大概率存在目标区域，最终输出候选区域。Anchor的首次出现是在Faster R-CNN[1]中，一组3×3的Anchor框是人为地按照3种长宽比(0.5,1,2)，3种缩放比(8,16,32)设置成，故并不适用于所有数据集。直到YOLO（You Only Look Once）的出现，在YOLO v2[4]中提出使用K-means聚类来自动生成一组适用于特定数据集的特定尺寸的anchor，简化工作量的同时，可以使网络的检测效果更好。

深度学习目标检测识别算法飞速发展的成果也呈现在生活垃圾检测识别领域：Ying Wang等人[5]采用RESNET作为特征提取网络，通过在城市自然环境中融合垃圾图像的方式扩充得到更多的数据集样本，从而提高了模型的检测精度和准确率。彭昕昀[6]采用利用数据增强对SSD算法进行垃圾识别检测，但SSD的检测框准确率较低。Abeywickrama等人[7]将垃圾分类视为图像分类，结合支持向量机与卷积神经网络对垃圾进行了识别分类，得到较高的识别率。宁凯等人[8]对YOLO v3[9]网络进行轻量化改进后在自制垃圾数据集上检测准确率较高，但是基于Darknet框架的YOLOv3在环境配置上较为复杂。

为了简化生活垃圾目标识别算法的网络结构，选取了轻量化的YOLO v5[10]算法，并且为提高其在生活垃圾上的检测精度，在识别外观相似物体时拥有更高的的泛化性，本文提出了基于YOLO v5与层次化分类结合的生活垃圾识别算法。首先将部分相似类别的垃圾数据合并，并对类别较少的数据集进行数据扩充，在一定程度上缓解了数据分布不均衡的问题。对这些预处理数据使用VGG16[11]网络依次训练得到八个分类器，将其与YOLO v5网络模型的输出融合，当检测网络输出类别与分类器对应的类别相同时，根据检测模型输出的位置信息将图片进行裁剪，输入到对应的分类器中进行具体的分类，用于相似类别垃圾的具体分类。再将预测的类别提供给检测模型的损失函数训练，以达到提升准确率的目的，且此改进算法结构简单清晰，训练方便。

一、算法设计分析

（一）YOLO v5算法介绍。YOLO系列是非常经典的one-stage[12]目标检测算法，YOLO v1[13]解决了通过直接拟合坐标位置、宽高存在很明显的缺陷；YOLOv2改为偏移量预测：采用了anchor机制局部预测，但YOLO v2还是存在小目标预测不准确的问题；YOLO v3随着resnet[14]的出现网络更深，backbone改为Darknet-53[15],检测头改为多尺度；YOLO v4[16]在bakbone更改为采用CSPDarknet53，数据增强采用了cutmix and mosaic,class label smoothing，网络结构采用了SPP、PAN、SAM网络。

YOLO v5在YOLO v4的基础上由Ultralytics公司提出，输入端依然采用Mosaic的数据增强方式，即将四张图片进行随机裁剪，再拼接到一张图上作为训练数据，提出了新的自适应聚类设定anchor大小，将预测框和真实框的groundtruth进行比对，计算两者差距后反向更新迭代网络参数，同时能够自适应检测图片的大小缩放。如图1所示，在Backbone上的主要是采用了Focus结构，主干网络采用CSP1_X结构，在Neck上采用了FPN结构和PAN结构。由于YOLO v5模型的目标检测功能还能够定位出垃圾所在位置，为我们生活垃圾识别的层次化分类算法研究提供了直接支持。

（二）基于YOLO v5与层次化分类算法的生活垃圾识别算法。YOLO v5极大地提升了算法的灵活性和速度的同时，依然具有较高的识别准确率。但是在实际实验过程中发现，该算法在检测外观相似物体的过程中，类别的判定效果不佳，这极大的限制了垃圾目标识别准确率的提升。为了缓解这一不足，我们提出了基于YOLO v5与层次化分类算法的生活垃圾识别算法，该算法任务的框架如图1所示。

图1 改进算法框架图

算法核心是将原YOLO v5网络输出的三个不同尺寸的特征提取面，对特征提取面上的每个点构建anchor，Anchor的构建是通过K-mean聚类以及对应本层特征提取面的尺寸得到的，再通过IOU和非极大值抑制筛选anchor，训练得到预测框和类别，根据预测框的位置对归并类所指的位置区域进行裁剪，将裁剪的图片送入已训练好的VGG16网络进行精确分类，将得到的类别重新赋予YOLOv5的类别输出，完成改进网络预测类别的替换。经过类别损失函数优化，进行迭代训练，计算最终准确率。

（三）损失函数。首先公式化定义基于YOLO v5与层次化分类算法的生活垃圾识别任务。由于希望控制在检测过程中的时间成本，需要对分类模型传入数据进行训练，以得到预训练模型参数。假设N表示样本数，概率分布p为期望输出，概率分布q为实际输出，则分类任务可以公式化为：

其中L(p,q)表示交叉熵损失函数，M表示类别数，p(xij)表示变量如相同即为1，否则为0；q(xij)表示样本i属于类别j的预测概率。

而检测识别部分class probability score采用了BCELoss（交叉熵损失），objectness score采用了BECWithLogitsLoss（BCEloss和sigmoid融合），bounding box采用了GIOU Loss：

其中A表示预测框，B表示标签框，C表示能够把A与B围起来的最小框，IoU表示预测框与真实框的交并比，C(A⋃B)表示C的面积减去(A⋃B)的面积。

改进后的算法损失函数如式（6）所示：

其中Lcls(ci,)是基于原算法改进的交叉熵损失函数，表示检测类别损失，其中表示真实类别，ci表示由（1）得到的预测类别。

改进算法的损失函数主要在Lcls部分联合分类网络的输出做了相应的改进，简化成更简单的类别细分任务，同时提高类别分类准确率。

二、试验及结果分析

（一）实验数据与实验环境。本文试验数据集在VOC2007基础上扩充到44个类别包含19933张生活垃圾图片，其中训练集13953张、测试集5980张分别用于原模型的训练和测试。后对其中的18个类别根据图片的相似度合并为8个类别，制成总类别为34的数据集。通过对原数据标签格式转换、划分数据类别、数据增强等方法完成数据预处理，缓解因数据不平衡对模型训练的影响。

本次实验操作系统为Ubuntu18，深度学习框架为Pytorch1.4，显卡为 NVIDIA Quadro P5000，加速环境为CUDA11。

（二）改进模型训练。本次实验分类模型的训练采用VGG16网络，根据数据标签的位置信息，将指定的目标图片对应位置裁剪得到分类训练数据，对类别较少的数据采用数据增强进行图片扩充，图片预处理后尺寸统一设置为90×90后输入网络得到8个预训练模型。检测模型部分采用最小的YOLO v5s预训练模型，原始图片尺寸统一设置为416×416，先后训练出原始44个类别的数据的检测模型以及相似类别合并后改进的34个类别的数据的检测模型。图2为改进模型进行3500次迭代的类别损失函数图像。

图2 改进的YOLO v5损失函数图像

训练参数betch size为32，共迭代次数为3500次，由图中可以看出，改进模型的损失函数在迭代到3000次左右时收敛到0.01附近，在迭代次数达到3500次后完成训练，得到最终检测识别模型。

（三）实验测试及结果分析。训练得到的改进模型通过在测试集上测试，得到的检测效果如图3所示。

图3 多场景下检测识别结果

图4中显示为多个类别在各种场景下的检测识别结果，效果不错，训练后的模型基本能够准确检测出目标位置，识别出其类别信息，且置信度较高。图4中在光线不平衡的环境下依然能够获得不错的检测识别结果，对于角落小物体的识别依然较为出色。说明改进模型在应对复杂环境下的生活垃圾检测识别具有较高的普适性。

图4 光线复杂环境下的检测识别结果

为评估YOLO v5与层次化分类算法对生活垃圾识别的有效性，本文采用平均准确率均值mAP（mean Average Precision）来作为衡量改进的垃圾目标识别模型性能的指标。mAP是指所有类别目标的准确率AP（Average Precision）的均值。对改进模型与原模型多次测试结果比较，绘制出表1。

表1 网络模型改进前后测试数据结果比较

由表1可以得出，原模型的AP88为33.1%，在垃圾目标识别情景下取得了不错的识别准确率。改进后的模型的AP88比原模型高出了0.5%，比YOLO v3高出了26.2%。此外改进模型的AP89、AP90较各模型均有较大程度的提高，且占据的优势越来越大，说明改进的模型在预测框与真实框具有更大的重叠区域时，表现得比原模型更为出色，准确率也更高，其出现误识别的几率更小。

三、结语

针对复杂环境下较多类别的垃圾检测识别，本文采用了基于YOLO v5与层次化分类算法，包括对数据预处理、算法原理分析、检测和分类模型的训练及后期优化。通过大量的实验对比发现，改进后的YOLO v5算法能够有效提升生活垃圾识别的准确率，在复杂环境下依然能够较原模型发挥出较大的优势。由于条件限制，本文所使用的数据集数量相对欠缺，对最终模型的表现关联较大，导致小物体的特征学习较为困难，在测试数据集上的泛化性有所欠缺。接下来的工作中，还需要对数据样本进行扩充优化分类模型，在保障现有性能的基础上解决更低IoU的识别准确率问题，进一步提高本算法的鲁棒性。