基于轻量化YOLOv5s的花椒簇检测研究
2023-06-15李光明弓皓斌袁凯
李光明 弓皓斌 袁凯
摘要:针对现有花椒簇检测算法模型参数量多、计算量大、检测速度低、很难部署到嵌入式设备的问题,提出一种基于轻量化YOLOv5s的花椒簇检测算法模型。首先将ShuffleNet v2主干网络替代原YOLOv5s中的主干网络进行重构;同时将SPPF嵌入至ShuffleNet v2骨干中;其次引入轻量级注意力机制CBAM;最后使用SIoU_Loss代替CIoU_Loss作为回归损失函数。试验结果表明:改进后的轻量化YOLOv5s网络参数降低85.6%,计算量降低87.7%,对花椒簇的检测精度mAP@0.5达到92.6%,较原YOLOv5s模型提高3.4%,mAP@0.5:0.95达到61.4%,检测时间为11 ms,相比原模型16 ms缩短31.3%,可以满足在现场环境下对花椒簇的检测。
关键词:花椒簇;YOLOv5s算法;轻量化;网络参数;检测精度
中图分类号:S225.93: TP391
文献标识码:A
文章编号:2095-5553 (2023) 04-0153-06
Abstract: In order to solve the problem that the detection algorithm model of pepper cluster has large number of parameters, large amount of computation, low detection speed and difficult deployment to embedded devices, a lightweight detection algorithm model of pepper cluster based on YOLOv5s was proposed. Firstly, ShuffleNet v2 backbone network was reconstructed instead of the original backbone network in YOLOv5s; At the same time, SPPF was embedded into the backbone of ShuffleNet v2; Secondly, lightweight attention mechanism CBAM was introduced; Finally, SIoU_Loss was used instead of CIoU_Loss as regression loss function to further improve the detection accuracy. The results showed that the network parameters of the improved lightweight YOLOv5s model were decreased by 85.6%, the computational amount was decreased by 87.7%, and the detection accuracy of mAP@0.5 for pepper clusters was 92.6%, 3.4% higher than the original YOLOv5s model, and mAP@0.5: 0.95 was 61.4%, and detection time was 11 ms, 31.3% less than the original model of 16 ms, which can meet the requirements for the detection of pepper clusters in complex field environments.
Keywords: pepper cluster; YOLOv5s algorithm; lightweight; network parameters; detection accuracy
0 引言
中國是世界上花椒产量最多的国家,2021年产量已达到50万吨以上。花椒产量的快速增长,人工采摘花椒速度过慢的问题越来越突出。所以研发全自动采摘花椒机器人具有很高的经济价值。花椒采摘机器人主要由花椒簇检测系统和机械臂控制系统组成,机械臂控制系统接收花椒簇检测系统所采集的信息对花椒进行采摘,因此对花椒簇进行快速、准确的检测是实现自动采摘[1-3]的关键。万芳新等[4]提出一种采用红色和绿色图像差异进行图像分割,并使用优化腐蚀滤除噪点和隔离带法除去非目标花椒串果实来实现花椒果实的识别和定位。汪杰等[5]提出一种HSV颜色空间与形状特征融合的方法对花椒图像进行分割。但由于采摘机器人所搭载的嵌入式设备片上资源有限[6],并且对检测速度要求很高,现有的复杂模型算法由于参数量过大,检测速度低,很难部署到嵌入式设备中,所以在保证检测精度的同时,减少模型参数量和计算量并且提高检测速度是花椒检测问题的难点。
近年来基于卷积神经网络(Convolutional Neural Networks, CNN)的目标检测算法因鲁棒性强、自适应性强及准确性高被广泛应用于图像识别等领域[7-8]。其可分为两种类型:第一种是“两阶段检测”算法,包括R-CNN[9],Fast-R-CNN[10],Faster-R-CNN[11]等标志性算法。另一种是一阶段的检测算法,主要包括YOLO系列[12-13],和SSD(Single Shot Multibox Detector)系列。但上述经典网络中均存在着模型过大和运算量过多等问题,很难进行嵌入式设备的移植。杨小冈等[14]提出一种在原始YOLOv5网络的基础上,使用MobileNetv3结构对模型主干网络进行重构,并添加特征提取网络MNtECA对网络模型进行压缩和加速的方法。张明路等[15]提出在YOLOv4算法主干网络中添加1×1卷积,用3×3的深度可分离卷积替代原有卷积减少网络参数量。张凡等[16]提出在YOLOv5s网络模型的基础上,结合EfficientNet网络中的MBConvBlock对主干网络进行重构,降低网络参数,并加入Stem模块,改进卷积注意力模块提升网络性能。本文针对已有算法网络参数量过大、检测效率低下、不易部署和难以移植到移动端等问题,基于YOLOv5s目标检测算法,提出一种轻量化花椒簇检测算法。
1 基于YOLOv5s的轻量化花椒簇检测算法
针对花椒簇检测,本文结合YOLOv5s与ShuffleNet v2的基础上提出一种轻量化的YOLOv5s网络模型,主要工作有以下4个方面。
1.1 ShuffleNet v2轻量化YOLOv5s主干网络
对比ShuffleNet v1[17]和MobileNet v2[18],发现4条轻量化网络结构的指导性原则:1×1卷积进行平衡输入和输出的通道大小;谨慎使用组卷积,注意分组数;避免网络碎片化;减少元素级运算。ShuffleNet v2模型使用3×3深度可分离卷积方法,计算量较3×3标准卷积方法减少88%左右,同时精度损失较小。深度可分离卷积包括深度卷积与1×1逐点卷积,深度卷积在每一个输入通道中都作用了一个固定尺寸卷积核,1×1卷积融合了深度卷积的输出特征,提高了网络的表达能力,同时没有改变输出特征图的尺寸。本文将轻量化分类网络ShuffleNet v2引进YOLOv5s目标检测算法主干中对花椒簇进行检测,为后续在嵌入式设备上部署提供可行方案。ShuffleNet v2模型基本结构如图1所示。
1.2 SPPF嵌入至shuffleNet v2骨干网络
快速空间金字塔池化(Spatial Pyramid Pooling-Fast,SPPF)模块是对空间金字塔池化SPP的简化,SPPF结构如图2所示,仅使用池化核为5的二维最大池化代替SPP结构中池化核分别为5、9、13的二维最大池化,减少计算量;同时采用并行的残差结构,将池化后得到的特征一部分通过Concat,一部分继续池化,通过池化操作得到不同层次的特征,提高模型对输入图像的特征提取能力。使用ShuffleNet v2轻量化网络后,模型的检测精度会在一定程度上变差,因此将SPPF嵌入至ShuffleNet v2骨干中,在几乎不增加模型大小的前提下来提高模型的检测精度。
1.3 轻量级注意力CBAM
本文引入轻量级注意力机制CBAM,更注重大区域内重要小目标信息。CBAM是一种轻量型模块,包括通道注意力模块(Channel Attention Module,CAM)与空间注意力模块 (Spatial Attention Module,SAM)两个子模块,这两个子模块分别在通道维度关注图片所输入的重要信息,在空间维度关注图片所输入的位置信息。改进后的YOLOv5s轻量化网络结构如图3所示,图中Shuffle-Block为ShuffleNet v2结构。
2 花椒簇数据集的制作与标注
本文以成熟的花椒簇为研究对象,在手机实地拍摄收集图像的基础上,以网络上的花椒图像作为补充,共获得500张照片。为提高网络模型的检测精度和泛化能力,在不同距离、不同光照、不同位置条件下对花椒进行实地拍照。拍摄时间为2022年8月,拍摄地点在陕西省白水县某种植园。对收集到的500张图片进行镜像翻转、旋转、改变亮度、对比度以及裁剪等方式进行扩充,共得到1 400张图片。本文使用LabelImg工具对收集到的花椒图像进行标注,标注后得到YOLO数据格式的.txt文件。1 400张图片按照7∶1∶2的比例划分为训练集、验证集及测试集。图片使用LableImg工具标注图如图4所示。
3 试验结果与分析
3.1 试验环境及参数设置
本文试验所用硬件设备如表1所示。软件环境为Python、Pytorch、CUDA、Pychram等。网络训练时批次大小(batch size)设置为32,动量因子(Momentum)设置为0.937,初始学习率(learning rate)设置为0.01,衰减系数设置为0.000 5,共训练300轮。
3.2 试验分析
试验中网络模型评价指标为平均精度均值mAP(mean Average-Precision)、平均每张图片检测时间Times、网络模型的参数量Parameters和模型计算量。IoU阈值取0.5时的mAP记为mAP@0.5,阈值为0.5~0.95之中取10个值得到的平均精度记为mAP@0.5:0.95。
试验1验证ShuffleNet v2主干网络对花椒簇检测的影响,在其他参数不变的情况下,在自建花椒簇数据集上进行消融试验,将主干网络改为ShuffleNet v2,添加SPPF模块,再添加不同轻量级注意力模块CBAM、ECA、SE进行试验,试验结果如表2所示。
从表2可以看出,当YOLOv5s主干網络改为ShuffleNet v2,添加SPPF模块和CBAM模块时,在模型大小与检测精度之间达到最佳的平衡,在mAP@0.5和mAP@0.5:0.95损失不大的情况下,其网络参数量为1 014 872,相比原模型7 053 910降低了85.6%;模型计算量为2.0 GFLOPs,相比原模型16.3 GFLOPs降低了87.7%;检测时间为11 ms,相比原模型16 ms缩短了31.3%。
试验2验证SIoU_Loss损失函数对花椒簇检测精度的有效性,在ShuffleNet v2、SPPF模块和CBAM模块应用的基础上,使用SIoU_Loss代替CIoU_Loss作为回归损失函数。试验结果如表3所示。结果表明本文算法对花椒簇的检测精度mAP@0.5达到92.6%,较原YOLOv5s模型提高3.4%,mAP@0.5:0.95达到61.4%。
试验3为分析本文网络模型的性能,与目前先进的目标检测网络EfficientDet和Faster-R-CNN在花椒簇数据集上进行比较试验。其余模型均按本文模型训练方式进行训练。试验结果如表4所示。结果表明本文花椒簇检测模型在损失精度很小的情况下,检测时间大大缩短,达到了预期目的。检测效果图如图5所示。
4 结论
1) 针对现有算法网络参数量过大、检测效率低下、不易部署到嵌入式等问题,基于YOLOv5s目标检测算法,本文提出一种改进的轻量化YOLOv5s花椒簇检测算法,利用ShuffleNet v2作为骨干网络替代原YOLOv5s算法中的骨干网络,将SPPF嵌入至ShuffleNet v2骨干网络最后一层,并引入轻量级注意力机制CBAM,最后将损失函数CIoU_Loss改进为SIoU_Loss,得到本文的轻量化检测花椒簇算法。
2) 改进后的轻量化花椒检测算法在保证原有检测精度不变的同时,有效地降低了模型计算复杂度,并提高检测速度。本文算法与原YOLOv5s算法相比,在自制花椒数据集上mAP@0.5提高3.4%,mAP@0.5:0.95只降低了1.3%,但参数量降低85.6%,计算量降低了87.7%,检测时间缩短31.3%,达到了预期效果。
3) 本文将进一步研究如何减少因轻量化所造成的精度损失,提高花椒检测精度,并研究如何在远距离的情况下识别树叶遮挡情况下的花椒簇,更进一步的完善花椒采摘机器人的实际应用。
参 考 文 献
[1] Kapach K, Barnea E, Mairon R, et al. Computer vision for fruit harvesting robots-state of the art and challenges ahead [J]. International Journal of Computational Vision & Robotics, 2012, 3(1/2): 4-34.
[2] 周俊, 刘锐, 张高阳. 基于立体视觉的水果采摘机器人系统设计[J]. 农业机械学报, 2010, 41(6): 158-162.
Zhou Jun, Liu Rui, Zhang Gaoyang. Design of fruit harvesting robot based on stereo vision [J]. Transactions of the Chinese Society for Agricultural Machinery, 2010, 41(6): 158-162.
[3] 阮承治, 赵德安, 陈旭, 等. 双指型农业机器人抓取球形果蔬的控制器设计[J]. 中国农机化学报, 2019, 40(11): 169-175.
Ruan Chengzhi, Zhao Dean, Chen Xu, et al. Controller design for realizing double-finger agricultural robot to grasp spherical fruits and vegetables[J]. Journal of Chinese Agricultural Mechanization, 2019, 40(11): 169-175.
[4] 万芳新, 白明昌, 贺志洋, 等. 自然场景下花椒果实的识别[J]. 中国农机化学报, 2016, 37(10): 115-119.
Wan Fangxin, Bai Mingchang, He Zhiyang, et al. Identification of Chinese prickly ash under the natural scenes [J]. Journal of Chinese Agricultural Mechanization, 2016, 37(10): 115-119.
[5] 汪杰, 陈曼龙, 李奎, 等. 基于HSV与形状特征融合的花椒图像识别[J]. 中国农机化学报, 2021, 42(10): 180-185.
Wang Jie, Chen Manlong, Li Kui, et al. Prickly ash image recognition based on HSV and shape feature fusion [J]. Journal of Chinese Agricultural Mechanization, 2021,42(10): 180-185.
[6] 王卓, 王健, 王枭雄, 等. 基于改进YOLO v4的自然环境苹果轻量级检测方法[J]. 农业机械学报, 2022, 53(8): 294-302.
Wang Zhuo, Wang Jian, Wang Xiaoxiong, et al. Lightweight real-time apple detection method based on improved YOLO v4 [J]. Transactions of the Chinese Society for Agricultural Machinery, 2022, 53(8): 294-302.
[7] 景亮, 王瑞, 劉慧, 等. 基于双目相机与改进YOLOv3算法的果园行人检测与定位[J]. 农业机械学报,2020, 51(9): 34-39, 25.
Jing Liang, Wang Rui, Liu Hui, et al. Orchard pedestrian detection and location based on binocular camera and improved YOLOv3 algorithm[J]. Transactions of the Chinese Society for Agricultural Machinery, 2020, 51 (9): 34-39, 25.
[8] 王丹丹, 宋怀波, 何东健. 苹果采摘机器人视觉系统研究进展[J]. 农业工程学报, 2017, 33(10): 59-69.
Wang Dandan, Song Huaibo, He Dongjian. Research advance on vision system of apple picking robot[J]. Transactions of the Chinese Society of Agricultural Engineering, 2017, 33(10): 59-69.
[9] Girshick R, Donahue J, Darrell T, et al. Rich feature hierarchies for accurate object detection and semantic segmentation [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2014: 580-587.
[10] Girshick R. Faster R-CNN [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2015: 1440-1448.
[11] Ren S Q, He K M, Girshick R, et al. Faster R-CNN: Towards real-time object detection with region proposal networks [J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017, 39(6): 1137-1149.
[12] Redmon J, Divvala S, Girshick R, et al. You only look once: unified, real-time object detection [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2016: 779-788.
[13] Redmon J, Farhadi A. YOLO9000: Better, faster, stronger [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2017: 6517-6525.
[14] 楊小冈, 高凡, 卢瑞涛, 等. 基于改进YOLOv5的轻量化航空目标检测方法[J]. 信息与控制, 2022, 51(3): 361-368.
Yang Xiaogang, Gao Fan, Lu Ruitao, et al. Lightweight aerial object detection method based on improved YOLOv5 [J]. Information and Control, 2022, 51(3): 361-368.
[15] 张明路, 郭策, 吕晓玲, 等. 改进的轻量化YOLOv4用于电子元器件检测[J]. 电子测量与仪器学报, 2021, 35(10): 17-23.
Zhang Minglu, Guo Ce, Lü Xiaoling, et al. Improved lightweight YOLOv4 for electronic components detection [J]. Journal of Electronic Measurement and Instrumentation, 2021, 35(10): 17-23.
[16] 张凡, 张鹏超, 王磊, 等. 基于YOLOv5s的轻量化朱鹮检测算法研究[J]. 西安交通大学学报, 2023(1): 1-12.
Zhang Fan, Zhang Pengchao, Wang Lei, et al. Research on lightweight Crested Ibis detection algorithm based on YOLOv5s [J]. Journal of Xian Jiaotong University, 2023(1): 1-12.
[17] Zhang X, Zhou X Y, Lin M X, et al. ShuffleNet: An extremely efficient convolutional neural network for mobile devices [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018: 6848-6856.
[18] Sandler M, Howard A, Zhu M L, et al. MobileNetv2: inverted residuals and linear bottlenecks [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2018: 4510-4520.
[19] Rezatohighi H, Tsoi N, Gwak J Y, et al. Generalized Intersection over Union: A metric and a loss for bounding box regression [C]. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019: 658-666.
[20] Zheng Z H, Wang P, Liu W, et al. Distance-IoU loss: Faster and better learning for bounding box regression [C]. Proceedings of the AAAI Conference on Artificial Intelligence, 2020, 34(7): 12993-13000.