垃圾是放错地方的资源, 未经分类的垃圾直接掩埋或焚烧, 造成的危害是难以想象的. 我国早在2017年就提出了《生活垃圾分类制度实施方案》^[1], 将垃圾分为了4大类进行分类处理, 以减轻环境污染. 在垃圾分类过程中, 我们可以在垃圾产生源头和在垃圾中端处理厂对垃圾进行分类, 但由于现阶段居民垃圾分类意识不强, 源头垃圾不能得到有效分类, 故需要在垃圾分选站对错分的垃圾进行二次分类. 目前中端垃圾分选站多依赖于人工分选, 在分拣精度和速度上都不够高, 同时恶劣的环境对工作者的身心健康影响巨大, 故分选自动化是垃圾分选站未来的发展趋势^[2].

针对中端垃圾分选的实际场景, 本文设计了一套适用于垃圾分选站的自动化分选系统. 首先通过工业摄像机拍摄传送带上的垃圾, 将视频传送到检测设备上; 然后通过OpenCV自带的接口读取视频中的每一帧图片, 将图片传递给经Tengine推理框架部署过的检测模型, 模型进行推理得到目标检测结果; 最后检测设备将检测结果中的物品类别和位置信息下发至机械臂, 机械臂进行垃圾夹取操作, 具体操作流程如图1所示. 现有的图像传输和机械臂控制技术足以满足上述要求, 故在自动化分选中最需要解决的是垃圾种类和位置的检测.

随着深度学习的高速发展, 各类目标检测算法被应用在垃圾检测中. 董霄霄等人^[3]基于YOLOv3提出了一种RP-YOLO的垃圾检测算法, 该算法使用RepVGG网络去替换YOLOv3的骨干网络, 并加入空洞卷积, 实现了复杂场景下的垃圾检测. 耿丽婷等人^[4]以SSD网络为基础, 更换了轻量级的主干网络, 并引入SK模块, 使网络在保证检测精度的前提下具有较快的检测速度. 魏铖磊等人^[5]在YOLOv3中引入多尺度感受野模块, 提高了网络的检测精度, 并更改了主干网络, 使得最终的网络能够对密集垃圾有较好的检测效果.

本文使用基于深度学习的目标检测算法实现垃圾的精准检测. 将目标检测算法运用在实际场景中, 最重要的就是要保证算法的检测精度和速度与实际场景相匹配. 由于垃圾分选站恶劣的工作环境和设备连续工作的机制, 传统的PC电脑无法适应在这种环境中, 故算法需要移植到嵌入式设备中. 但由于嵌入式设备算力较为低下, 故检测算法需在保证检测精度的前提下具备较快的推理速度. 本文以YOLOv4目标检测网络为基础, 搭建了一个新的检测网络, 该网络可以更好地移植到算力和内存较低的嵌入式设备上.

1 背景知识 1.1 YOLOv4算法简介

YOLOv4网络在YOLOv3^[6]的基础上对其主干网络、颈部网络和损失函数等多方面进行改进, 在保证检测速度的前提下, 极大地提升了检测精度, 同时对小目标问题有较好的检测能力. YOLOv4网络主要由3大部分组成, 分别是进行特征提取的CSPDarknet53主干网络, 进行特征融合的颈部网络SPP网络和PAN网络, 以及用于目标检测的检测头网络, 其具体的网络结构如图2所示. 但在垃圾检测中, YOLOv4^[7]因其网络参数量较大、检测速度较慢, 不适合移植在嵌入式设备上, 本文预对YOLOv4网络进行改进, 在保证网络精度的前提下, 尽可能压缩网络参数量、提升检测速度.

1.2 MobileNetV3简介

MobileNetV3^[8]是一个轻量级网络, 适合部署在嵌入式设备和算力有限的移动设备上. MobileNetV3由1×1普通卷积和深度可分离卷积堆叠形成的倒残差结构搭建而成, 其部分倒残差结构中添加SE注意力机制以提升网络的特征提取能力, 具体的网络结构如图3所示. 本文将以MobileNetV3为基础, 重新设计YOLOv4中的主干网络, 使得检测网络更为轻量化.

1.3 Ghost模块简介

Ghost模块^[9]由华为在2020年提出, 其本质是在特征提取过程中, 首先使用1×1卷积对输入特征通道数进行压缩, 然后利用逐层卷积对其进行特征提取, 最后将前两步得到的特征层进行堆叠, 完成特征提取, 其特征提取过程如图4和图5所示. 实验发现这种特征提取方式在避免提取重复特征的同时, 大大降低了计算量, 加快了模型的推理速度.

2 本文工作

在YOLOv4中, 原始的骨干网络CSPDarknet53为一个中量级的网络, 其大量使用3×3普通卷积, 导致参数量和网络架构较为庞大, 并不适合部署在资源有限的嵌入式设备上. 而MobileNetV3^[10]中使用深度可分离卷积代替传统的卷积模块, 在保证特征提取精度的前提下, 减少了模型的参数量, 提升了模型的推理速度. 基于以上, 本文首先将YOLOv4中的CSPDarknet53网络替换为MobileNetV3网络; 同时对MobileNetV3在网络层级上进行调整, 使得主干网络能更好地匹配YOLOv4检测网络^[11]; 接着设计了一种全新的bottleneck结构替换掉层级调整后的MobileNetV3中的bottleneck结构; 然后在主干网络的第3、4、5层级上添加大残差边结构; 最后在颈部网络之前添加CA注意力机制模块, 构建更适合于嵌入式设备的检测网络, 网络总体框图如图6所示.

2.1 改进的主干网络 2.1.1 层级均衡的MobileNetV3网络

MobileNetV3网络主要由5大层组成, 分别对应提取到的5种不同的feature map. 对416×416的输入图片, YOLOv4使用3个检测头对其进行检测, 故本文选取MobileNetV3的后3层作为特征提取的输出. 输出feature map为52×52和13×13的第3层和第5层仅有2个bottleneck进行特征提取, 会造成图片特征提取不均衡, 影响检测精度, 故本文为这两层各增加3个5×5的bottleneck, 达到图片特征提取均衡的效果, 具体网络结构如图7所示.

2.1.2 基于Ghost模块的G-Mobile bottleneck

MobileNetV3中的深度可分离卷积相较于传统卷积在模型参数量和推理速度上有了很大的提升^[12], 但研究发现深度可分离卷积在特征提取过程中会出现特征重复, 重复的特征可以由其余特征经过线性变换得到. 针对上述研究, 发现通过少量的卷积和堆叠的方式可得到相较于深度可分离卷积参数量更小, 速度更快的Ghost模块^[13], 其计算量仅为深度可分离卷积模块的 $1/s$ , 具体论证结果如式(1)所示:

$ \begin{gathered} {{{r}}_s} = \frac{{n \times {h{'}} \times {w{'}} \times c \times k \times k}}{{n/s \times {h^{'}} \times {w{'}} \times c \times k \times k + \left( {s - 1} \right) \times n/s \times {h{'}} \times {w{'}} \times d \times d}} \\ = \frac{{c \times k \times k}}{{1/s \times c \times k \times k + \left( {s - 1} \right) \times {\left( {s - 1} \right)}/s \times d \times d}} \approx \frac{{c \times s}}{{s + c - 1}} \approx s \\ \end{gathered} $

(1)

其中, 分子为普通卷积的复杂度, 分母为Ghost模块的复杂度, $n \times {h{'}} \times {w{'}}$ 为输出的特征尺寸, $ s $ 代表每个通道产生的总映射, $ c $ 为输入特征图且数值一般较大, $ n/s $ 指经过普通卷积输出的特征图数, $ k \times k $ 为卷积核的大小, $ d \times d $ 为线性操作核的大小, 与 $ k \times k $ 大小相似.

研究表明^[14], MobileNetV3中纺锤形结构并不有利于梯度回传, 会使得网络提早出现过拟合的现象. 针对上述结论, MobileNeXt网络构建了一种两头大中间小的沙漏型结构来代替MobileNet系列网络中使用的两头小中间大的纺锤形结构, 并将深度可分离卷积层放置在bottleneck两头, 使其有更多的信息可以从底层传递到顶层, 有利于梯度回传, 同时可以提取到多种空间信息, 具体结构如图8所示.

本文综合Ghost模块中的特征提取思想和MobileNeXt网络中的bottleneck结构设计思想, 搭建了一种全新的bottleneck结构, 用以替换层级调整后的MobileNetV3中的bottleneck结构, 具体模型结构如图9所示, 按照特征尺寸是否需要压缩可分为步长为1和2的两种bottleneck结构. 本文设计的bottleneck融合了Ghost模块的思想性, 采用少量逐层卷积和堆叠的方式作为特征提取的模块, 同时将MobileNetV3中先大后小的纺锤形bottleneck, 设计成先小后大的沙漏型结构, 在一个bottleneck中采用两个特征提取模块, 以提取更多的空间特征信息. 采用本文设计的bottleneck可以有效缓解MobileNetV3中纺锤形网络结构带来的信息损失和梯度混淆, 同时可以进一步压缩模型计算量, 使得模型精度和速度得到进一步的提升.

2.1.3 添加大残差边的主干网络

CSP结构^[15]的本质是一个大的残差网络, 使用在主干网络的每一大层上. 在主干网络上增加CSP结构, 可以在防止模型过拟合的基础上, 有效地保留前一层提取到的图像特征, 从而进一步提升网络的检测精度. 本文借鉴了CSPNet网络的思想, 在主干网络的后3层添加了大残差边结构, 为使得残差边输出feature map数和该层的feature map数相同, 该结构在跳跃分支上添加了一个1×1卷积, 以保证网络拼接的平衡性. 添加大残差边的目的是在尽量不加深轻量级G-MobileNet网络深度的前提下, 获取更多的样本特征, 提高主干网络的特征提取能力, 进而提升检测网络的检测精度. 未在网络前两层添加大残差边结构的原因是前两层主干网络中bottleneck较少, 且每一个bottleneck都自带一个残差网络, 再添加大残差边结构会出现冗余的现象, 且会导致网络参数量增加. 以G-MobileNetV3网络第3层为例, 改进后的GR-MobileNetV3网络如图10所示.

2.2 改进的颈部网络

在YOLOv4中, 颈部网络的主要作用是通过反复上采样和下采样对提取到的特征信息进行融合, 使得检测网络的每个检测头都可以得到不同层级的图像信息, 因此在特征融合之前获得更为准确的特征信息尤为重要, 注意力机制作为一种即插即用的网络模块, 其主要作用是使得网络在训练过程中更能注意到需要的特征信息.

CA注意力机制^[16]是一种通道注意力机制, CA注意力机制将通道注意力沿W和H分解为两个1维特征编码, 分别沿2个空间方向捕获长度依赖关系和位置信息. 两种信息互补的应用于特征图像, 可以增强网络对需要特征的学习^[17]. 同时CA注意力机制相较于其余的通道注意力机制, 几乎没有额外的运算成本的开销, 很适合运用在轻量级网络上. 其具体模型图如图11所示.

本文将CA注意力机制添加在主干网络输出的特征提取层之后, 颈部网络之前, 得到GR-MobileNetV3-CAYOLOv4检测网络, 使得网络忽略训练图片中的背景干扰信息, 更多的注意到图片中需要被注意到的信息.

2.3 二分K-means算法优化anchor box

在YOLOv4算法中, 对于每个检测头需要3个初始的anchor box, 故需要在数据集中选出9个不同的anchor box. 对于不同的数据集来说, anchor box也不尽相同. 在YOLOv4中, 我们经常使用K-means聚类算法, 通过输入数据集的标定位置, 聚类生成9个不同的anchor box. K-means算法的本质是随机生成9个不同的anchor box, 通过数据集中各个标记框距anchor box中心的距离将其余标记框划分为9簇; 然后再找出每簇标记框最中心的anchor box, 作为新的9个anchor box; 接着重新计算中心距离, 划分为新的9簇; 如此循环往复, 直到生成的9个不同的anchor box位置不在变换为止, 即可认为聚类完成.

K-means算法过分依赖于初始聚类中心的选择, 如果初始聚类中心选择不好, 可能会导致聚类生成的anchor box不是最优解而是局部最优解, 针对K-means算法这一缺陷, 本文使用二分K-means算法^[18]对anchor box进行求解.

在聚类中, 我们认为一个簇中的各个元素靠得越近, 则该簇的聚类越为准确, 二分K-means算法则采用了这种思想. 二分K-means算法首先将整个数据集作为一个簇, 将该簇进行K为2的K-means聚类, 即将数据集分为两个簇; 接着计算2个新簇中各个样本到中心anchor box距离的算术平方和, 选择算术平方和较大的那个簇类使用K为2的K-means聚类算法再次划分; 然后计算现有的3个新簇中各个样本到中心anchor box距离的算术平方和, 接着进行二分K-means操作; 最后将整个数据集划分为9个簇, 选择9个簇中的中心anchor box作为目标检测算法的最终anchor box. 二分K-means算法解决了K-means算法过分依赖初值的缺陷, 不断地使用K为2的K-means操作, 使得聚类划分更为科学, 算法不会陷入局部极值, 从而提升检测算法的检测精度.

3 实验分析

本文实验使用Windows 10操作系统的主机; 内存为16 GB; GPU为NVIDIA GeForce RTX 3060; 使用Python 3.8; 学习框架为PyTorch 1.71+cuda 11.0. 实验中通过对比主干网络的分类能力验证本文主干网络的性能, 通过将本文设计的检测网络同YOLOv4和YOLOv4-tiny等算法进行对比, 验证本文网络的检测能力.

3.1 数据集建立

本文对检测网络的总体结构进行了重新设计, 所以YOLOv4中提供的预训练权重不能直接进行使用, 故本文须通过分类训练获得主干网络的预训练权重. 本文使用华为提供的垃圾分类数据集对主干网络构成的分类网络进行训练, 该数据集中包含40种垃圾图片, 图片总数为17 000余张. 通过分类网络训练后, 可以得到分类网络的垃圾分类识别率和相应主干网络的预训练权重.

在垃圾检测网络的训练中, 本文将标注好的华为垃圾分类数据集以及自行扩充的垃圾检测数据集进行结合, 得到了本文使用的垃圾检测数据集, 该数据集中包含43种垃圾图片, 图片总数为23 000余张.

3.2 数据集划分

根据本文数据集样本数量和检测难度, 可将数据集分为简单样本、中等样本和困难样本3种类别. 其中简单样本指样本数量较多, 类内差异较小、类间差异较大的样本; 困难样本指样本数量较少, 类内差异较大、类间差异较小的样本; 其余为中等样本. 图12所示的茶叶渣和食用油桶, 其形状颜色固定, 较容易检测, 故被定义为简单样本; 而图13的饮料瓶, 因其颜色种类较多, 且掺杂有形变, 较难进行检测, 故被定义为困难样本. 按照上述定义可将本文数据集进行统一划分, 得到图14所示的3类样本占比图, 由图可知本文使用的数据集困难样本占比较大.

3.3 模型训练

本文实验首先使用垃圾分类数据集训练分类网络, 然后剔除分类网络中的全连接层和分类器权重, 得到主干网络的预训练权重, 将预训练权重添加到检测网络中, 使用垃圾检测数据集进行检测网络的训练.

分类网络训练中, 采用统一的训练学习率0.000 1, 模型输入分辨率统一转换成(416, 416, 3), 统一训练400代. 检测网络训练中, 采用自适应学习率, 模型输入分辨率统一转换为(416, 416, 3), 冻结和不冻结主干网络各训练100代.

3.4 评价指标

针对主干网络的分类能力, 本文选取测试集上分类效果最好的5代权重的识别率, 求取平均值作为分类网络的平均识别率(mean recognition, MR), 具体如式(2)所示, 以衡量主干网络的分类能力, 使用模型参数量Params来衡量模型的大小.

$ MR = 0.2 \times \sum\limits_{i = 1}^5 {\frac{{CS\left( i \right)}}{{TS\left( i \right)}}} $

(2)

其中, ${{CS}}\left( i \right)$ 和 ${{TS}}\left( i \right)$ 分别为测试集中第 $i$ 代分类正确的样本数和总样本数.

本文以(mean average precision, mAP)作为模型检测精度的评价指标, 以每秒钟图片的传输帧数(frames per second, FPS)和参数量Params作为检测速度的评价指标^[19].

3.5 垃圾分类网络实验

本文使用垃圾分类数据集对各主干网络构造的分类网络进行训练, 具体结果如表1所示.

从表1可以发现, CSPDarknet53网络参数量大, 具有较好的分类能力, 但推理速度较慢; 层级均衡的MobileNetV3参数量降低了80%, 但其检测精度同样下降明显; G-MobileNetV3网络的模型参数量进一步降低, 同时检测精度有所提升, 但精度与CSPDarknet53仍有一定差距; 本文使用的GR-MobileNetV3网络, 模型参数量虽然略有增加, 但精度却得到了很大的提升, 相较于CSPDarknet53网络精度提升了2.1%, 故本文将YOLOv4中的主干网络替换为GR-MobileNetV3网络.

3.6 垃圾检测网络实验 3.6.1 检测网络消融实验

对检测网络进行消融实验, 第1组为MobilenetV3-YOLOv4网络, 其余组别依次添加√标注的网络结构, 可以得到表2所示的实验结果.

组别2层级调整后的MobilenetV3-YOLOv4网络的推理速度相较于MobilenetV3-YOLOv4略微下降, 但mAP值增加2.7%; 将bottleneck 更换为本文设计的G-Mobile bottleneck, 组别3的FPS和mAP均得到了提升; 接着在主干网络的第3、4、5层添加大残差边结构, 组别4网络检测精度提升3.7%, 但推理速度略有下降; 组别5是添加CA注意力机制的本文网络, 其网络运行速度上几乎没有下降, 但mAP提升0.7%; 组别6是在本文设计的检测网络基础上, 使用二分K-means聚类算法进行聚类操作, 相比于使用K-means算法在检测精度上提升了0.3%. 综合分析表2中数据可以发现, 本文采用的网络以较小的检测速度, 换取了9.1%的检测精度提升.

3.6.2 检测网络对比实验

对比不同检测网络在本文数据集上的训练效果, 可得到如表3所示实验结果.

分析表3, 由Params指标可以看出本文构建的网络与YOLOv4-tiny网络处于同一量级, 但mAP值明显优于YOLOv4-tiny网络. 本文网络与YOLOv4和YOLOv3相比, 参数量约为它们的1/10, 且FPS提升了51%和31%, mAP值和YOLOv4网络相当, 比YOLOv3网络提升了10.7%. 综合实验结果可以看出, 本文构建的网络在保证检测精度的前提下, 能够极大的提高检测速度, 可以部署在算力和内存较低的垃圾抓取嵌入式设备上.

3.6.3 不同类别样本检测能力对比实验

根据样本数量和检测难度, 数据集被分为简单样本、中等样本和困难样本3种类别, 将这3种类别的测试样本图片输入到各检测网络中, 可得到不同网络下各类样本的mAP值, 此处提取典型的简单样本、中等样本和困难样本的mAP值, 如表4所示. 其中茶叶渣、食用油桶、锅是典型的简单样本, 毛绒玩具、包、剩饭剩菜是典型的中等样本, 插头电线、饮料瓶、金属食品罐头是典型的困难样本.

分析表4发现, 本文网络对于茶叶渣、食用油桶和锅等简单样本的检测能力较高, mAP值可以达到65%以上, 对于插头电线、金属食品罐等困难样本的检测能力较弱, mAP值仅为10%左右. 综合数据集样本数据发现, 检测效果不好的多为样本数据量不足, 或类内差异大的困难样本, 如本文数据集中的插头电线、金属食品罐, 它们的训练图片数量仅为150余张, 饮料瓶因其类内差异过大, 且无大量数据集样本支撑, 因此检测效果均不理想.

3.6.4 检测网络实例对比

我们使用5种不同种类的垃圾图片输入到YOLOv4、YOLOv4-tiny和本文构建的检测网络中, 可以得到如图15的检测结果.

对比图15的检测结果, 可以发现YOLOv4-tiny在第3、4两张图片上未检测到物体, 第2张图片将茶叶错误识别为菜叶, 第1、5张图片检测框定位不准确, 虽然其FPS很高, 但不适用于垃圾检测. 对比YOLOv4和本文检测网络可以发现, 本文检测网络在第3张图的检测效果低于YOLOv4网络, 其余图片中检测能力均高于YOLOv4网络, 尤其是在第2张这种有遮挡的图片中, 本文网络能同时检测出茶叶和陶瓷器皿两种物体, 并且定位准确, 说明本文的检测网络有较强的抗遮挡能力, 同时具备更快的推理速度, 更适用于算力较差的嵌入式设备上, 满足现阶段垃圾检测的要求.

4 结束语

基于垃圾检测现状, 本文对YOLOv4进行改进, 设计了一种适用于垃圾检测的网络. 相较于YOLOv4网络, 本文网络的参数量减少了89%, FPS提升了51%, 同时具有和YOLOv4相似的mAP值. 通过对大量图片进行实验, 发现本文检测网络在垃圾相互遮挡方面具有较好的检测效果, 适合部署在算力和内存较低的嵌入式设备上. 针对数据集, 后续研究中可通过增加训练样本的数量来提高网络的检测精度; 同时针对检测网络的颈部网络进行进一步的改进, 以提升网络的检测精度.