基于改进YOLOv5的番茄果实识别估产方法
作者: 杨健,杨啸治,熊串,刘力
摘 要:为了实现大棚环境中番茄的智能在线产量估算,提出了一种基于改进的YOLOv5(You Only Look Once v5)番茄识别算法,对自然生长状态下的番茄果实产量进行统计和估算。首先,使用可分离视觉转换器(Separable Vision Transformer,SepViT)替换骨干网络的最后一层,以增强骨干网络与全局信息之间的联系并提取番茄特征;其次,引入WIOU(Wise Intersection over Union)损失函数和Mish激活函数,以提高收敛速度和精度。试验结果表明,改进后的检测模型在平均精度(mean Average Precision,mAP)方面达到了99.5%,相较传统的YOLOv5模型提高了1.1个百分点,每张图像的处理时间为15 ms。此外,改进后的YOLOv5算法对密集和遮挡情况下的番茄果实识别效果更好。
关键词:大棚番茄;YOLOv5;注意力机制;损失函数
中图分类号:S641.2 文献标志码:A 文章编号:1673-2871(2024)06-061-08
An improved YOLOv5-based method for tomato fruit identification and yield estimation
YANG Jian, YANG Xiaozhi, XIONG Chuan, LIU Li
(School of Mechanical and Electrical Engineering, Chengdu University of Technology, Chengdu 610059, Sichuan, China)
Abstract: In pursuit of intelligent real-time yield estimation for tomatoes in greenhouse environments, we introduce an enhanced YOLOv5 tomato recognition algorithm aimed at the statistical assessment and estimation of tomato fruit yield in their natural growth conditions. Our approach involved two key enhancements: firstly, we substituted the final layer of the backbone network with a Separable Vision Transformer to augment the connectivity between the backbone network and global context, thereby facilitating tomato feature extraction; secondly, we incorporated the WIOU loss function and employ the Mish activation function to enhance convergence speed and accuracy. Experimental findings demonstrate that the improved detection model achieves a mAP score of 99.5%, reflecting 1.1 percentage points enhancement compared to the conventional YOLOv5 model, and the processing time for every image is 15ms. Furthermore, the improved YOLOv5 algorithm exhibits superior recognition performance for densely populated and occluded tomato fruits.
Key words: Greenhouse tomato; YOLOv5; Attention mechanism; Loss function
随着我国经济的发展,农业正在向现代化转型,温室大棚已成为现代农业的重要支撑,而大棚番茄种植也在全国范围内快速发展。准确估算大棚番茄的产量是种植规划和营销策略的重要组成部分,实时统计和预测番茄果实产量信息,并加以相应的生产控制,可以解决番茄产能波动和生产过程中的不连续性等问题。然而,大棚内的番茄茎叶和果实密集且相互重叠[1],限制了果实图像特征的识别,进而影响了番茄产量的准确估算。因此,实现准确获取番茄果实视觉信息是支持智能估产的重要前提。
传统的番茄识别方法通常是手动设计特征提取器来提取目标特征,如梯度直方图(Histogram of Oriented Gradient,HOG)、尺度不变特征变换(Scale-invariant feature transform,SIFT)和Haar-like特征等,这些特征再输入支持向量机(Support Vector Machine,SVM)、AdaBoost、随机森林(Random Forest,RF)等分类器进行分类和识别[2]。Liu等[3]应用伪色去除技术和SVM达到了94.41%的识别率,但在番茄遮挡和重叠等复杂情况下,效果不佳。李寒等[4]采用RGB-D相机拍摄图像,通过预处理得到水果轮廓,采用K均值聚类与自组织映射算法进行番茄识别,识别率达到了87.2%,但轮廓提取易受光照影响。Malik等[5]采用改进的HSV颜色空间算法和分水岭分割算法分离红色番茄,但识别率只有81.6%。而马翠花等[6]采用基于密度的局部自适应阈值分割算法提高番茄检测精度,对未成熟番茄的识别率只达到77.6%。这些传统方法泛化性差、计算复杂度高,难以提取出合理的特征。此外,它们往往没有考虑到大棚环境中的复杂影响因素,对各种特征的变化鲁棒性不足,难以满足实际需求。而近年来兴起的深度卷积神经网络为获取番茄果实视觉信息提供了更好的方法。
随着深度学习的不断发展,目前的目标检测算法可以分为两类。第一类是基于候选区域的two stage算法,例如R-CNN[7]、Fast R-CNN[8]和Faster R-CNN[9]等。这类算法通过两个步骤来进行计算:首先选择候选框,然后对候选框进行分类或回归。这种方法具有高鲁棒性和低识别错误率,但运算时间长,占用磁盘空间大,并且对图像信息进行了重复计算,不适合实时检测应用。第二类是基于网络模型的one stage算法,例如YOLO[10-11]和SSD[12]等。这类算法通过在图像上以不同尺度进行遍历抽样,并利用卷积神经网络提取特征后直接进行回归。这种方法的网络模型较为简单,识别速度快,特别适合实时检测应用[13]。其中,YOLO模型是Redmon等[10]于2017年提出的对象检测模型,与two stage算法相比,YOLO模型使用单个前馈网络直接预测边界框及其对应的类,作为一种基于回归的检测器,不仅支持快速实时识别,其精度也能够满足一般应用要求。
2018年,周云成等[14]提出了一种基于Fast R-CNN的双卷积链算法,对番茄果实识别的平均精度为63.99%。2021年,张文静等[15]将AlexNet作为Faster R-CNN的骨干网络,在50%的重叠度下,番茄检测的平均精度为83.9%,每幅图像的处理时间为245 ms。2020年,刘芳等[16]通过改进YOLO网络模型,可在复杂环境中有效地检测番茄果实,识别准确率达到97.13%。同年,Liu等[17]提出一种用于检测番茄果实的改进YOLOv3模型,在轻度遮挡条件下,番茄果实的识别准确率达到94.58%,但在遮挡严重时表现较差。2021年,成伟等[18]使用广义交并比(Generalized Intersection over Union[19],GIOU),作为改进YOLOv3算法的位置损失函数,在番茄果实检测方面的平均精度达到95.7%,每幅图像的平均处理时间为15 ms[19]。
综上所述,笔者针对大棚环境下番茄果实重叠、枝叶遮挡、光照不均等复杂情况,提出一种改进的YOLOv5目标检测算法。首先,使用可分离视觉转换器(SepViT)替换骨干网络的最后一层,以增强骨干网络与全局信息的联系和提取番茄的特征;其次,利用WIOU损失函数和Mish激活函数来同时提高收敛速度和精度,使得YOLOv5能够在大棚环境下获得对番茄果实较高的检测精度。
1 番茄图像数据集制作
在笔者研究中使用的番茄数据集部分是在四川省成都市拍摄的,部分是从参考文献中获得的。这些图像有效地反映了番茄生长的复杂性,如大小和光线的差异、番茄之间的遮挡以及枝叶的遮挡。在各种条件下,包括侧光、背光、阴影等,随机拍摄了1000余张番茄图像,图1所示为已建立的番茄数据集,拍摄时间为2022年7月,该番茄识别数据集已经在Zenodo上公开发布,可以在https://zenodo.org/record/8429227上获取。
为了增强原始图像数据,提高数据集的普适性和鲁棒性,笔者采用了空间级和像素级数据增强方法,且同时使用了这两类增强方法,如图2所示。空间级变换技术修改边界框和图像,而像素级变换技术修改图像,同时保持边界框不变。
2 基于改进YOLOv5的番茄识别模型
2.1 YOLOv5目标检测算法
YOLOv5是一种基于卷积神经网络(Convolutional Neural Networks,CNN)的无锚点对象检测算法,相比其他对象检测算法具有更高的效率和准确性,其原理如图3所示。
YOLOv5算法的基本原理是将输入的图片分成S×S个栅格化单元格。如果检测到的目标中心落入特定的单元格中,则该单元格负责检测目标,即自身有目标的概率为[Pobj=1]。然后,每个单元格预设生成B个先验边界框,每个边界框与真实值边界框的交并比为IOU。因此,图像中的目标位置和类别预测可以表示为[S×S×B×(4+1+C)]的张量,其中4表示先验边界框的坐标[(x,y)]和宽高[(w,h)],1表示置信度得分,总共有5个特征参数,C表示使用的数据集中目标的类别数。通过连续回归训练真实边界框,可以得到最终预测的目标位置、置信度和类别信息。最后,保留置信度最高的边界框,以筛选出最佳的识别结果[20]。
2.2 深度可分离视觉注意力机制
最近几年,视觉转换器(Vision Transformer,ViT)在各种计算机视觉任务中取得了巨大的成功,其性能在关键领域甚至超过了传统的卷积神经网络[21-22]。然而,这些性能提升通常是以增加计算复杂度和参数数量为代价的,而深度可分离视觉注意力机制可以通过平衡计算成本和准确性来提高模型的性能[23]。
笔者将骨干网络的最后一层替换为SepViT Block,以增强模型的特征提取能力,并优化与全局信息的联系。在SepViT Block中,Depthwise selfattention (DWA)和Pointwise selfattention(PWA)可以减少计算量,并实现窗口内的局部信息通信和全局信息交互。首先,输入的特征图划分为多个窗口,每个窗口被认为是特征图的一个输入通道,不同类型的信息被包含在每个窗口中;然后,对每个窗口及其像素标记执行DWA,这有助于融合通道的空间信息,类似于MobileNet[24-26]中的深度卷积。DWA的操作如下:
[DWA(f)=Attentionf∙WQ,f∙WK,f∙WV] 。
其中[f]是特征标记,由窗口标记和像素标记组成,[WQ]、[WK]和[WV]代表常规自注意力任务中的query、key、value三个线性层,PWA类似于MobileNet中的pointwise convolution操作,不同之处在于pointwise convolution用于融合不同通道的信息,而PWA则是在窗口之间建立连接。完成DWA操作后,PWA建立窗口之间的关系,并通过Layer Normalization (LN)和Gelu激活函数生成注意力图。PWA的运行过程如下: