基于改进YOLOX的茶叶嫩芽目标检测研究

Research on Tea Sprouts Object Detection Based on Improved YOLOX

摘要: 茶产业是我国进出口贸易商品的一大重要方面，茶在我国有着悠久的文化底蕴，与我国人民的生活密切相关。为了满足名优茶的茶叶嫩芽智能采摘需求，本文首先建立了自然环境下的茶叶嫩芽数据集，并提出了一种基于Swin-Transformer的改进YOLOX茶叶嫩芽检测模型——YOLOX-ST。该模型将Swin-Transformer作为原始YOLOX模型的骨干网络，提高了模型整体的检测精度，并引入了CBAM注意力机制，解决复杂环境背景下容易错检漏检的情况。实验结果表明，该模型的mAP值达到了79.12%，比原始模型提高了5.2%，精确度达到了90.45%，比原始模型提高了4.62%。与同系列的YOLOv3、YOLOv4、以及YOLOv5模型相比，YOLOX-ST的mAP以及准确率最高分别提升了7.09%和6.43%，拥有良好的检测精度与模型泛化能力。由此可见，该模型为茶叶嫩芽的智能化采摘奠定了一个良好的基础。

Abstract: Tea industry is an important aspect of China’s import and export trade commodities. Tea has a long cultural heritage in China, and is closely related to the life of our people. In order to meet the acquirement of premium tea sprouts, this paper first established the dataset of tea sprouts based on natural environment, and proposed a modified YOLOX tea sprout detection model based on Swin-Transformer—YOLOX-ST. The proposed model used the Swin-Transformer as the backbone network, which improves the overall detection accuracy of the model. And it also introduced the CBAM attention mechanism to solve the problem of miss-detection and wrong detection in the complex environment. Experimental results showed that the proposed model has a mAP value of 79.12%, which is 5.2% higher than the original YOLOX model, and the accuracy has achieved 90.45%, which is 4.62% higher than the YOLOX model. Compared with YOLOv3, YOLOv4, and YOLOv5, the mAP and accuracy rate of YOLOX-ST model increased by 7.09% and 6.43% at most, respectively, which has good detection accuracy and model generalization ability. This model has laid a good foundation for the intelligent picking of premium tea sprouts.

1. 引言

茶产业作为我国一大进出口的重要产业，在各大农业商品中占据着重要地位。名优茶拥有着良好的茶叶品质以及优良的加工工艺，从古至今都受到人们的追捧。为了满足我国茶产业高品质需求的快速发展，需要不断提高生产高端名优茶的生产能力，对茶叶的采摘质量、采摘速度与加工工艺方面提出更多、更高的要求 [1]。目前，茶叶嫩芽的采摘技术主要分为人工采摘及机械采摘。机械采摘，是一种使用机械采茶机的方法，可以较好地实现茶叶嫩芽的快速采摘，但机械采摘的方法缺乏对茶叶嫩芽的选择性，所得到的茶叶嫩芽存在杂质，如茶杆、茶病叶等，增加了后期对茶叶嫩芽加工筛选的工作量，以至于严重影响了名优茶叶的质量。人工采摘可以实现对茶叶嫩芽的精确采摘，但其效率低下，容易花费大量额外的人工成本，且当今社会人口老龄化严重，不符合名优茶产业的快速、高品质发展的需求 [2]。因此，研究茶叶嫩芽的智能化采摘迫在眉睫。自然环境下，茶叶嫩芽背景复杂，且生长姿态各异，因此，茶叶嫩芽的智能化采摘研究中，首要任务是研究茶叶嫩芽的识别与检测。

茶叶嫩芽的检测方法，大部分依赖于茶叶嫩芽的整体形状，以及嫩芽在颜色方面的特征提取，如吕军 [3] 等人提出了基于区域亮度自适应校正的茶叶嫩芽检测模型，对高亮度图像进行2 × 3分块和局部区域亮度自适应校正后，获得了较高的召回率。杨福增 [4] 等人提出了基于颜色和形状特征的茶叶嫩芽识别方法，通过对RGB颜色空间中提取茶叶图像的G分量，采用双阈值方法对图像进行分割，最后分辨出了茶叶嫩芽。但上述方法往往容易出现过分割的问题，且自然环境下的茶叶嫩芽背景复杂，仅仅依靠颜色及形状无法达到精确识别的要求。姜苗苗 [5] 等人提出了基于颜色因子与图像融合的茶叶嫩芽图像检测算法，实现了茶叶嫩芽与老叶的区分，但该方法容易因图像噪声等因素影响，且准确率仍有待提高。因此，由此可知，传统图像处理的方法已经不能满足高精度的茶叶嫩芽识别与检测需求。

近年来，深度学习一直是诸多学者研究的方向，基于深度学习的茶叶嫩芽检测方法也相继出现。施莹莹 [6] 等人创新性地使用了多种光照条件、多种类型的茶叶嫩芽数据集作为YOLOv3的网络输入，使得模型具有较强泛化能力与检测效果。邹倩 [7] 等人提出一种改进YOLOV3的茶叶嫩芽检测方法，模型检测精度达到了79.9%，相比于YOLOV3模型，精度提高了7.8%。

上述研究虽然都取得了不错的效果，但在模型检测精度上仍有待提高，且使用的模型较为老旧，不利于后续的发展。本文基于自然环境下建立的茶叶嫩芽数据集，以YOLOX为基础模型，将Swin-Transformer模块引入到其骨干网络当中，并且增加了CBAM注意力机制，提高了模型整体泛化能力及复杂背景下对茶叶嫩芽的检测能力。

2. 图像采集与预处理

2.1. 图像采集

目前网络上还没有公开的茶叶嫩芽数据集，因此，本文构建了一个自然环境下的茶叶嫩芽数据集。于2021年3月至2021年4月，在贵阳市花溪区羊艾茶园进行茶叶嫩芽原始图像的采集，拍摄设备为iPhoneXR手机后置摄像头，像素为1200万，总共采集并挑选出6357张茶叶嫩芽原始图像，部分图像如图1所示。

Figure 1. Original tea sprouts images

图1. 茶叶嫩芽部分原始图像

2.2. 图像预处理

使用Labelimg软件对茶叶嫩芽原始图像进行标注工作，标注后的图像及标签文件如图2、图3所示。标注时仅将茶叶的嫩芽部分框入标签框中，并将标注的嫩芽类别命名为teabud。标注完成后，对所有图像进行带标签的翻转、镜像、对比度调整等操作，如图4所示。最后，本文的数据集从原始的6357张图像扩充到了10,693张图像。

3. 基于改进YOLOX的茶叶嫩芽识别检测模型

3.1. YOLOX原始模型

YOLOX是2021年Zheng Ge [8] 等人提出的一个检测模型，其结构如图5所示。由图可知，整个YOLOX可以分为三个主要的部分，分别为CSPDarknet、中间部分的FPN以及最后的YOLO Head部分。输入的茶叶嫩芽图片，首先会在CSPDarknet部分中进行对茶叶嫩芽的特征提取，随后输入到FPN部分，进行特征融合，最后在YOLO Head部分对特征点进行判断，判断特征点是否有物体与其对应。

Figure 2. Label of tea sprout

图2. 茶叶嫩芽标注

Figure 3. The XML file of tea sprout

图3. 茶叶嫩芽标签文件

CSPDarknet作为整体网络的主干，具有四个重要的特点。第一，使用了多个残差网络。CSPDarknet中的残差卷积由一个1 × 1的卷积和一个3 × 3的卷积组成，其残差边部分不做任何处理。第二，使用了Focus网络结构。Focus网络首先是在一张图片中每间隔一个像素取出一个值，此时获得了四个独立的特征层，然后将这四个特征层进行堆叠操作，拼接后的特征层相对于原先的3通道变成了12通道。第三，使用了SiLU激活函数。SiLU激活函数具有无上界但有下界、平滑、非单调的特性。它的效果要优于ReLU函数。第四，使用了SPP结构。在YOLOX中，SPP模块被用在了主干特征提取网络中，通过不同池化核大小的最大池化进行茶叶嫩芽的特征提取，提高整个网络的感受野。

Figure 4. Data expansion

图4. 数据扩充

YOLOX网络利用提取的多特征层来进行目标检测，一共提取三个特征层，分别位于主干部分CSPDarknet的中间层，中下层以及最底层，利用FPN特征金字塔，可以获得三个加强特征层，接着利用这三个shape的特征层传入YOLO Head来获得最终的预测结果。

Figure 5. Structure of YOLOX base model

图5. YOLOX原始模型结构

3.2. Swin-Transformer

Transformer [9] 是近年来深度学习中大火的一个模块，首先被提出来应用于NLP领域，并且在该领域大获成功。受此启发，2021年Alexey Dosovitskiy [10] 等人提出了Transformer在CV领域的应用——Vision Transformer (VIT)，文章中给出的最佳模型在ImageNet1K公共数据集上能够达到88.55%的准确率，也就是说，Transformer这模块应用在CV领域这一想法是可行的。

Swin-Transformer [11] 作为VIT的一个改进版本，与VIT不同的是，Swin-Transformer使用了类似卷积神经网络中的层次化构建方法(Hierarchical feature maps)，比如特征图尺寸中有对图像进行下采样4倍，8倍以及16倍，这样的骨干网络有助于在此基础上构建目标检测，实例分割等任务。它不仅具有Transformer模块关注全局信息的能力 [12]，而且采用了移动窗口的方法，实现了跨窗口连接，使模块可以关注到相邻窗口的相关信息，这样可以扩大模型的感受野，提高模型整体的效率。

Swin-Transformer的结构图如图6所示。

Figure 6. Structure of Swin-Transformer

图6. Swin-Transformer结构图

Swin-Transformer首先将图片输入到Patch Partition模块中进行分块，即每4 × 4相邻的像素为一个Patch，然后再通过线性嵌入层(Linear Embedding)对每个像素的通道数据做线性变换。接着通过四个Stage构建不同大小的特征图，除了Stage 1中先通过一个线性嵌入层外，剩下三个部分都是先通过一个Patch Merging层进行下采样，最后对于分类网络，后面还会接上一个Layer Norm层、全局池化层以及全连接层得到最终输出。

3.3. CBAM注意力机制

CBAM (Convolutional Block Attention Module)是一种轻量的注意力模块 [13]，可以分别在通道和空间维度上进行提高注意力的操作，结合了两个方面的注意力信息。与单一的关注通道特征的注意力机制和单一关注空间注意力机制的方式相比，CBAM具有更加优异的性能。CBAM的结构如图7所示。

CBAM包含CAM (Channel Attention Module)和SAM (Spartial Attention Module)两个子模块，其结构图如图8所示。

CAM模块的主要作用是获取茶叶嫩芽特征中不同通道特征之间的相互关系 [14]，并学习相应的权重分布，从而找到并突出其中最有意义的茶叶嫩芽通道特征。对于输入的茶叶嫩芽特征 $F\in R^{C\times H\times W}$，CAM首先使用最大池化(Maxpool)和平均池化(Avgpool)的操作，将特征的空间信息聚合在一起，得到两个空间向量。然后将这两个空间向量分别送入多层感知机网络(MLP)，最后将输出的两个特征进行相加融合，通过Sigmoid函数激活后生成通道注意力权重 $M_c$。将 $M_c$ 乘以茶叶嫩芽特征 $F$，变得到中间特征 $F^{\prime }$。CAM的计算公式如式(1)所示。

Figure 7. Structure of CBAM Attention Mechanism

图7. CBAM注意力机制结构图

Figure 8. Structure of CAM and SAM

图8. CAM及SAM结构图

$M_c\left(F\right)=\sigma \left(W_1\left(W_0\left(F_{avg}^C\right)\right)+W_1\left(W_0\left(F_{\max }^C\right)\right)\right)$ (1)

式中 $\sigma$ 为Sigmoid激活函数， $\left(F_{\max }^C\right)$ 和 $\left(F_{avg}^C\right)$ 分别表示最大池化和平均池化， $W_0\in R^{C/R\times C}$ 和 $W_1\in R^{C/R\times C}$ 为MLP网络的两个权重值。

与CAM不同，SAM的主要功能是关注茶叶嫩芽特征的空间位置信息，给出其最有信息量的空间位置，是对CAM的一种补充机制 [15]。对于CAM输入的 $F^{\prime }$，通过Maxpool及Avgpool，获得两个单通道特征，并将这两个单通道特征进行拼接操作，最后使用3 × 3的卷积进行通道融合，经过Sigmoid函数运算，得到空间注意力权重 $M_s$，将 $M_s$ 乘以中间特征 $F^{\prime }$，最终得到所输出的特征 $F^{″}$。SAM的计算公式如式(2)所示。

$M_s\left(F\right)=\sigma \left(f^{3\times 3}\left(F_{\max }^S;F_{avg}^S\right)\right)$ (2)

式中 $\sigma$ 为Sigmoid激活函数， $f^{3\times 3}$ 表示3 × 3的标准卷积， $F_{\max }^S$ 和 $F_{avg}^S$ 分别表示最大池化和平均池化。

3.4. YOLOX-ST模型

在Swin-Transformer取出的三个有效特征层中，分别为下采样8倍、16倍以及32倍的特征层，正好对应原始YOLOX网络中CSPDarknet所输出的特征层。因此，本文提出了一种以Swin-Transformer作为骨干网络的YOLOX-ST模型，并在模型下采样后引入了CBAM注意力机制，以此来提升模型对茶叶嫩芽的检测精度。YOLOX-ST模型的结构图如图9所示。

Figure 9. Structure of YOLOX-ST model

图9. YOLOX-ST模型结构图

4. 实验结果及分析

4.1. 实验环境

实验系统环境使用的GPU是RTX3050显卡，使用的编程语言为Python，版本为3.8，训练模型所使用的CUDA版本为11.1，对应的PyTorch版本为1.7.0，如表1所示。

Table 1. Experimental environment configuration

表1. 实验环境配置

4.2. 实验结果分析

本文输入的图片大小为640 × 640，有三个通道。用此数据集分别对YOLOX原网络和改进后的网络YOLOX-ST进行训练和验证，所用图像共10,693张，训练集、验证集、测试集按8:1:1的比例来划分，epoch设置为400，学习率设置为0.0001，batch size设置为16，使用SiLU损失函数。

为了验证YOLOX-ST整体模型的有效性，本文将训练集分别送入YOLO系列不同网络模型中进行训练作为对比，所得到的准确率及mAP如表2所示。

Table 2. Comparison of different models

表2. 不同模型对比

^a加粗字体表示最优结果。

由表2可知，本文提出的YOLOX-ST模型不管是从mAP上还是从准确率上都有一个良好的提升。相较于原始的YOLOX模型，mAP和准确率分别提升了5.2%及4.62%，而与其他模型相比，优化后的模型mAP最高提升了7.09%，准确率最高提升了6.43%。

训练后的模型与原始YOLOX模型的准确率曲线及Loss曲线比较如图10、图11所示。由图可知，YOLOX-ST模型在准确率上的提升速度比原始YOLOX模型要快，在100 epoch之后便显现出了明显的差距。且YOLOX-ST模型的起始Loss也比原始模型的低两个点，下降速率也比原始模型较快，由此可以更直接地观测到本文所提出模型的有效性。

Figure 10. Precision curve

图10. 准确率曲线

Figure 11. Loss curve

图11. Loss曲线

为了验证Swin-Transformer以及CBAM注意力机制的合理性，本文进行了消融对比试验，结果如表3所示。

Table 3. Ablation test

表3. 消融试验

^b“√”表示在网络中使用该方法。

由表3可知，无论是Swin-Transformer还是CBAM注意力机制，应用在模型中时都保持着有效性，也即验证了本文提出的方法的合理及可行性。

5. 结论

本文针对自然环境下的茶叶嫩芽数据集，提出了一种基于YOLOX的改进模型YOLOX-ST，该模型引入了Swin-Transformer模块作为骨干网络，有效地优化了模型对于茶叶嫩芽的特征提取，并在茶叶嫩芽特征提取下采样后加入了CBAM注意力，加强了茶叶嫩芽特征的信息量，提升了模型整体的检测准确率。实验表明，与其他YOLO系列模型以及原始YOLOX模型相比，YOLOX-ST模型具有准确率高，泛化能力较好的特点，为后续的茶叶嫩芽智能化采摘奠定了一个良好的基础。