基于YOLO神经网络的垃圾检测与分类

张 伟，刘 娜，江 洋，李清都

(1.上海理工大学 光电信息与计算机工程学院，上海 200093； 2.上海理工大学 医疗器械与食品学院，上海 200093； 3.重庆市互联网信息办公室，重庆 401120)

随着经济的迅速发展，我国城镇人口数量急剧增加，城市生活垃圾也日益增多，而垃圾分类则是处理垃圾公害的最佳解决方案[1]。当下的关注热点在于将垃圾进行无害化、资源化处理，因此合理有效地对垃圾进行回收利用具有一定的意义。对垃圾进行有效地分类是处理垃圾回收问题的前提[2]。目前，针对人工分拣垃圾环境差、分拣效率低等问题，研究人员提出建立垃圾识别与分拣智能化和自动化系统，通过利用深度学习技术来降低人工分类成本，节约人力资源[3]。

垃圾检测属于计算机视觉领域的目标检测任务。传统的检测方法以HOG(Histogram of Gridie) 检测器[4]、基于滑动窗口的VJ(Viola-Jones)检测器和DPM(Deformable Part Model)检测器[5]为代表。卷积神经网络(Convolutional Neural Network，CNN)架构于2012年被首次提出，并逐渐成为了热门研究方向。经过多年发展，深度学习领域已涌现出VGGNet、GoogLeNet、ResNet等多种神经网络算法。最先被用于目标检测的卷积神经网络是R-CNN[6](Region-CNN)算法。该算法使用卷积神经网络进行特征提取，使用Bounding Box Regression修正目标边界框。不过R-CNN算法也有其弊端，串行式CNN前向传播耗时长，且CNN特征提取、SVM(Support Vector Machines)分类和边框修正3个模块是分开训练的，对于存储空间的消耗较大。为了克服R-CNN的耗时问题，Fast R-CNN算法直接采用一个CNN提取全图特征，减少了对存储空间的占用。R-CNN和Fast R-CNN算法存在相同的缺陷，即存在耗时的Selectice Search，因此这两种算法无法实现端到端的物体检测。为了解决这个问题，文献[7]提出了一种区域提议网络(Region Proposal Network， RPN)，该网络可以直接连接到Fast R-CNN 网络的后一个卷积层，以得到对应的检测对象框和分数信息，该网络也被称作为 Faster R-CNN。上述目标检测算法都是二阶段算法，针对二阶段目标检测算法普遍存在运算速度慢的缺点，研究人员提出了YOLO(You Only Look Once)[6]算法，该算法可在一个步骤中完成对物体的分类和定位。YOLO和SSD[8](Single Shot MultiBox Detector)这类一阶段算法仅使用一个CNN网络就能直接预测不同目标的类别与位置。SSD的网络结构基于VGG16构造，其检测速度较快。YOLO的网络结构类似于GoogLeNet，其先将图片尺寸放缩到固定尺寸，然后将图片传入CNN网络进行处理，用非极大值抑制进行边界框[9]的冗余裁剪，从而处理掉大批冗余，得到最后的预测结果。YOLO将物体检测作为一个回归问题进行求解，仅一次求解便可以得到图像中所有物体的位置、类别以及相应的置信概率。

本文探讨了CNN的精细特征对垃圾分类任务的适用性，针对不同实验场景，制作了特定的垃圾数据集。本文提出采用基于YOLO的神经网络模型来进行垃圾检测与分类，可有效解决垃圾分类的人工预处理问题。

1 垃圾分类检测流程

本文提出的垃圾检测与分类方法设计包括数据集制作以及模型调整两部分。垃圾检测与分类的方法设计整体流程如图1所示。

图1 设计流程图

本文实验使用的数据来源于相机拍摄，数据集共计20 000张图片。垃圾样本共计9类，分别是易拉罐、烟盒、烟头、果皮、包装袋、纸、塑料瓶、瓜子壳、树叶。YOLO模型相较于R-CNN等两阶段算法拥有一定的速度优势。相较于YOLO模型早期版本，YOLOv4模型的识别精度更高，且有利于边缘计算，对于项目落地有较强的可行性，因此本文选用YOLOv4模型进行研究。

2 垃圾分类检测方法设计

2.1 数据集的获取

由于项目机器人作业地点为郑州东站，故在郑州东站进行数据采集，采集时间为夏季，天气以晴天为主，少量雨天。因郑州东站环境较复杂，所以数据采集的背景选择非常重要。本实验数据采集的背景选择了30余处，基本覆盖所有场景，确保了实验数据的完整性和准确性。通过D435双目相机获取视频，视频分辨率为1 920×1 080，每秒传递帧数为30帧。为了获得最佳的训练数据，本文充分考虑了垃圾的各种形态，选择多角度环绕拍摄，避免垃圾之间产生重叠。对拍摄完的视频进行解析，每秒提取2张图片，格式为JPEG。为了保证图像的质量和清晰度，对解析后的图片进行人工筛选，然后对筛选后的图片使用labellmg工具进行标注。标注时，为了保证数据的准确性，需确保目标框为垃圾样本的最大外接矩形。对标注后的数据，进行人工验收，防止出现漏标误标。后期则根据模型测试的效果，对易误检的目标或背景再进行样本补充，提高识别的准确率。采集的垃圾样本如图2所示。

图2 垃圾样本

2.2 数据集的划分

对所有图片进行命名排序，根据标注产生的txt格式文件编写相应的代码使其生成训练集、验证集和测试集。

训练集是模型训练过程中使用的图像数据[10]。模型效果与训练集的选择有直接关系。训练集对初始模型拟合与参数优化有重要作用。本文的训练集分为单类垃圾数据和多类垃圾数据。

验证集被用于在每个Epoch完成之后进行模型状态和收敛情况的检验，它不参与梯度下降的过程，仅负责调整迭代次数和学习率等超参数。此过程不参与学习参数的确定。

测试集用于对训练完成的最终模型进行测试，通过测试来确定该模型检测与分类的准确率，确保该模型在未来实际应用中的正确性和有效性。

2.3 垃圾识别分类方法

本文探讨的垃圾检测与分类方法使用基于YOLOv4的网络作为实验训练模型，并将其用于全局和局部的特征提取。首先，YOLOv4网络主体结构如图3所示，主要分为3部分：主干网络、颈部网络、头部网络。主干网络采用CSPDarknet53，颈部网络是SPP(Spatial Pyramid Pooling)模块+PANet(Path Aggregation Network)，头部网络则采用YOLOv3。

图3 YOLOv4网络主体结构

图3中，CBM和CBL模块是融合了批量归一化并使用Mish和Leaky ReLU激活函数的卷积层；CSP模块由多个残差块叠加而成；SPP模块中的Concat操作则是将大小相同的特征图进行通道拼接。

为了在不增加训练成本的基础上增加模型的准确率，本文采用了数据增强和边界框回归损失函数。数据增强能增加输入图片的多样性，从而增强复杂环境下模型的鲁棒性[11]。传统的目标检测方法使用的是均方误差(Mean Squared Error，MSE)，与检测框评测指标IOU(Intersection over Union)存在误差，故引入了 LossIou，其定义如式(1)所示。

LossIou=1-IoU

(1)

注意力机制模块可以根据全局图像，从大量信息中筛选有用信息，抑制其它无用信息，增强特征图中重要的空间和通道特征，使网络在训练过程中可抓住目标特征重点学习。本文选用卷积注意力机制模块(Convolutional Block Attention Module，CBAM)，并将其嵌入到YOLOv4模型中。如图4所示，本文选择在每个检测头前嵌入一个CBAM模块。

图4 CBAM嵌入示意图

在CBAM模块中[6]，对于一个中间层的特征图F∈RC×H×W，C表示特征图的通道，H表示特征图的高度，W表示特征图宽度。一维的通道注意力特征图MC∈RC×1×1，二维的空间注意力特征图MS∈R1×H×W，整个过程为

F1=MC(F)⊗F，F2=MS(F1)⊗F1

(2)

式中，⊗表示逐元素相乘。先将通道注意力特征图与输入的特征图相乘得到F1；然后计算F1的空间注意力特征图；再将两者相乘得到输出F2。

通道注意力模块以及空间注意力模块的计算过程如下文所述。CBAM通道注意力模块使用最大池化和平均池化对特征图进行空间维度上的压缩，得到两个不同的通道背景描述，即Fmax和Favg。再通过由多层感知器(Multilayer Percep-tron，MLP)组成的共享网络对Fmax和Favg进行计算[12]，得到通道注意力特征图 ，如式(3)所示。

MC(F)=σ(MLP(AvgPool(F))+MLP(MaxPool(F)))=

σ(W1(W0(Favg))+W1(W0(Favg)))∈RC×1×1

(3)

式中，σ表示Sigmoid激活函数；MLP表示共享全连接层[12]；W0表示共享全连接层的第1层；W1表示共享全连接层的第2层；输出向量长度为C。

通道注意力模块(Channel Attention Module)如图5所示。

图5 CBAM通道注意力模块

CBAM空间注意力模块先在通道维度上进行最大池化和平均池化，得到两个不同的空间背景描述，再使用卷积操作[13]生成空间注意力特征图

MS(F)=σ(f7×7([AvgPool(F)；MaxPool(F)])) =

σ(f7×7([Favg；Fmax]))∈RH×W

(4)

式中，f7×7表示7×7的卷积层。

图6所示为空间注意力模块示意图(Spatial Attention Module)。

图6 CBAM空间注意力模块

相比于FPN(Feature Pyramid Networks)，PANet的效果更好，但计算量也更大。本文中YOLOv4特征集成采用PANet，故需先将PANet进行简化，将PANet的addition改成concatenation。PANet主要包含FPN、Bottom-up path augmentation、Adaptive feature pooling、Fully-connected fusion共4个部分[14]。FPN通过融合高低层特征来提升目标检测的效果，其对小尺寸目标检测效果的提升更为显著。浅层特征一般是边缘形状等特征，考虑到网络浅层特征信息对于实例分割较为重要，故引入Bottom-up path augmentation。Adaptive feature pooling用每个感兴趣区域提取不同层的特征来做融合。Fully-connected fusion则向原有分割支路引入一个前背景二分类的全连接支路[15]，通过融合这两条支路的输出来得到更精确的分割结果。

本文选用Mish函数作为激活函数，如图7所示，其正值无上界的特点避免了由于封顶而导致的饱和。另外，Mish函数保证了每一点的平滑性，从而获得了较好的梯度下降效果[16]。

图7 Mish函数示意图

Mish(x)=x×tanh(In(1+ex))

(5)

2.4 数据集训练部分

本次实验平台基于Ubuntu16.04系统，显卡为TITAN RTX 2080Ti，内存为32 GB。训练部分基于YOLOv4-custom.cfg对数据集进行模型训练，其中batch值为64，subdivision值为16。因本实验数据集共有9类，所以max_batches值被设置为18 000次。Steps=xxx，本文将xxx设置成max_batches的80%和90%。YOLOv4有3个YOLO层，将每个层里面的classes值设为数据集类型数目9。数据集按照二八原则进行划分，即20%为验证集，剩余80%作为训练集，并编写Python脚本文件来分离训练集和测试集，在Main中生成train.txt和val.txt文件。训练时打开Trainning的batch和subdivisions，将Testing部分注释掉。本实验基于YOLO预训练模型开始训练，当训练的loss趋于平稳时，停止训练。使用训练得到的模型对垃圾进行检测，对易误检的垃圾进行数据补充，从而继续提升模型精度。

3 实验结果与分析

针对本文提出的模型调整与算法改进措施，分别进行对比实验来检测性能。目前平均精度(Average Precision，AP)是多类别目标检测及其分类中最重要的评价指标，因此本次实验采用AP作为垃圾检测与分类的评价指标[17]。各类别垃圾在YOLOv4以及改进的模型上的精确度表现如表1所示。

表1 不同垃圾类别的AP比较

由表1可知，采用本文提出的算法模型后，有8类垃圾的检测精确度高于YOLOv4的精确度，其中对PlasticBottle的检测精确度提升最高，提升了6.36%；识别效果最好的类别则为Cans，识别精度可达94.56%。本实验还对比了CSPResNEXt50和CSPDarknet53在垃圾检测与分类方面的表现，具体结果如表2所示。

表2 CSPResNEXt53与CSPDarket53的参数对比

由表2可知，CSPResNeXt50仅包含16个卷积层，其感受野为425×425，包含20.6 MB参数； CSPDarkNet53包含29个卷积层，725×725的感受野，共有27.6 MB参数。因此，虽然CSPResNEXt50在分类方面的表现优于CSPDarkNet53，但其在检测方面的表现则不如CSPDarkNet53[18]。

此外，本实验还对比了使用K-means聚类算法优化Anchor box后和未经优化的YOLOv4以及本文引入注意力机制和K-means的YOLOv4模型在本文数据集上进行训练后的性能变化，具体结果如表3所示。

表3 不同模型的效能评估

由表3可知，当IOU阈值设置为0.5时，相较于原始YOLOv4，使用K-means聚类算法的YOLOv4-1的map提升了1.08%，使用K-means聚类算法并引入注意力模块的模型的map提升最为明显，提升了2.81%。改进YOLO模型对垃圾的识别结果如图8所示。

图8 改进后的算法识别效果

4 结束语

本文针对现实生活中人工分拣垃圾效率低、任务重且环境差等问题，基于卷积神经网络的模型进行改进，提出了进行垃圾分类与检测的新方法。该方法能够在一定程度上改善垃圾分类耗时耗力的问题，并准确检测垃圾，完成分类。本文在YOLO模型中融入了K-means聚类算法，并向该模型中加入CBAM注意力机制，研究了注意力机制对垃圾检测算法性能的影响，并选择经过优化的网络架构来完成垃圾检测与分类任务。下一步将对模型进行压缩，尝试构造更好的轻量级网络结构并将其用于模型训练，从而提升网络模型的实时性。