基于改进YOLOv5s的大米外观品质无损检测模型

顾海洋, 蒋雪松, 沈 飞, 郁鑫智, 周宏平

(南京林业大学机械电子工程学院1,南京 210000)

(南京财经大学食品科学与工程学院2,南京 210000)

水稻经过砻谷加工、碾米和抛光过程,会产生不同程度的破损,影响了大米的产量和外观品质。另外,大米在加工和储存过程中很可能出现破损和发霉的现象,而表型缺陷是种子品质评价的重要指标。在大米贸易流通中,外观品质是影响大米价格的重要因素之一,因此提高大米分选质量具有重要意义[1]。

随着深度学习在目标检测中的不断发展,研究者将卷积神经网络应用到农产品检测当中,取得了较好的效果,可以为大米品质检测提供参考。

基于神经网络的目标识别技术,大致可以分为2类,一类是一阶段目标检测算法(one-stage),包括YOLO[2]、SSD[3]、EfficientNet、Retina Net[4]等;另一类是二阶段(two-stage)目标检测算法,包括R-CNN[5]、Fast R-CNN[6]、PVANet、SPPNet[7]和Mask R-CNN等。

目标检测网络在农作物品质检测中应用也越来越广泛[8-11]。王巧华等[12]提出了基于YOLOv5s[13]网络的改进识别网络,用来检测棉花种子破损,检测准确率为95.4%。范晓飞等[14]利用以MobileNetV1为骨架的YOLOv5s网络对玉米种子进行识别,其对玉米种子检测精度为93%。顾伟等[15]利用Resnet50网络代替SSD算法中的VGG网络,该网络对群体脱绒棉籽中的破损棉籽的检测精度达到96.1%。

YOLOv5s模型是在YOLO4基础上改进而来,具有检测速度快、精度高和泛化能力强等优点。为进一步提高大米目标的识别精度,本研究对YOLOv5s网络进行改进,添加小目标检测层以提高小目标识别率,为使模型轻量化,用Ghost模块代替YOLOv5s的骨干网络,再添加CBAM注意力机制加强对大米目标重要特征提取。利用改进后的YOLOv5s网络对大米进行识别,以提高大米识别率。

1 基于CG-YOLOv5的大米检测模型

1.1 YOLOv5s检测模型

YOLOv5算法是在2020年提出的,有YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x共4个模型。YOLOv5s网络输入图片大小为640(像素)×640(像素),经过一系列卷积相关操作得到32倍下采样的20×20特征图、16倍下采样的40×40特征图、8倍下采样的80×80特征图。20×20、40×40、80×80的特征图对应检测大、中、小目标。

1.2 模型的改进

1.2.1 小目标检测层的添加

YOLOv5s原始模型只有3个检测层,负责检测小目标的是80×80的特征图,该层特征图包含的语义信息相对较小,易出现小目标漏检现象。为加强对小目标的识别率和准确率,增加一个小目标检测层[16],即160×160的特征图。同时采用多尺度特征图相融合的方式,将浅层特征图与深层特征图融合,来保留更多小目标语义信息。图1是添加检测层的网络结构图。修改的小目标检测层有2个方面。

图1 添加检测层的网络结构图

原80×80的特征图经过上采样与骨干网络中的第一个CSP1_1的特征图相融合,再通过CSP2_1模块和卷积操作得到160×160的特征图,即小目标检测层尺度。

160×160的特征图经过卷积模块与原80×80的特征图相融合,再经过CSP2_1模块和卷积操作得到新的80×80的特征图。

1.2.2 Ghost模块

使用Ghost模块替换YOLOv5s骨干网络中的CSP1_模块,Ghost模块卷机原理如图2所示。第1个普通卷积用来降低通道数量,减少网络参数量,第2个卷积是采用深度可分离卷积中的深度卷积,后将2个卷积进行拼接操作。当输入向量为H(行数)×W(列数)×N(通道数)时,输出向量为H×W×M(M为偶数)时,输入向量经过M/2个1×1×N的卷积核将其压缩为H×W×M/2的向量。后经过深度卷积,深度卷积对输入向量的每1个通道进行1个不同的k×k×1的卷积运算,经过M/2个深度卷积得到M/2个H×W的向量。

图2 Ghost模块卷积原理图

Ghost模块可以有效减少参数量。例如,对于1个N通道的输入向量,先经过M/2个1×1×N的普通卷积核,后经过M/2个3×3×M/2的普通卷积核,这个过程所需要的计算量为:

(1)

如果经过1个Ghost模块完成相同的工作,需要的参数量为:

(2)

在Ghost模块运行时,因为输入通道和输出通道数相同,即N=M,所以2种情况下的比值为:

(3)

由此可知,当输入向量通道数很大时,2种方法的卷积方式参数量的比值为0.18。

1.2.3 CBAM特征增强模块

CBAM注意力机制是一种结合空间和通道的注意力机制的模块[17],包含通道注意力机制和空间注意力机制。CBAM注意力机制结构如图3所示,特征图先经过通道注意力机制模块,后传入空间注意力机制模块,完成特征的提取。

图3 CBAM注意力机制结构

输入特征图,先经过通道注意力机制,分别进行自适应平均池化和自适应最大池化,得到2个1×1通道权重矩阵,在经过类似瓶颈层的结构减少参数量,2个特征在向量相加后经过sigmoid激活函数,生成单通道特征图。单通道特征图与输入特征进行乘法运算,生成中间特征图。经过空间注意力机制,通过最大池化和平均池化操作得到2个二维向量拼接,经过一个卷积层和sigmoid激活函数后与中间特征图相乘得到输出特征图。

1.3 CG-YOLOv5的大米检测模型

为提高大米识别的精度,提出了CG-YOLOv5大米检测模型,网络结构图如表1所示。对模型的改进为:添加一个小目标检测层来提升大米检测精度;在特征提取中,将CBAM注意力机制和CSP1_模块进行融合,组成新的模块CBAMC3,从而加强对大米的特征提取;将YOLOv5s骨干网络中的CSP1_模块替换成Ghost模块,能有效减少参数量,提高检测速度。

表1 网络结构表

2 数据集的制作和模型训练

2.1 数据集的获取

为了识别大米,需要采集相关图像建立图像数据集。以江苏兴华大米为研究对象,图像采集装置为USB工业高清相机,型号为MS500D,采集图片的格式为jpg。

实验在传送带上实时采集样本图片。将大米分为3类,分别为正常米、垩白米和碎米。各类大米图像如图4所示。训练模型采用VOC数据格式,用Labelimg进行文件标注,将外观良好的大米标注为Normal,将白色不透明和有缺陷的大米标注为Chalky,将断裂的大米标注为Broken。

图4 大米分类图片

实验共采集大米图片2 000张,通过丰富数据集的数量,可以有效解决因数据集不足导致网络拟合较差的问题,从而提高网络的泛化能力。数据集进行网络训练时,将训练和验证、测试集的比例设为8∶1∶1。经统计,训练集中正常米10 702粒,垩白米9 312粒,碎米7 283粒。测试集中用于评估模型性能的正常米447粒,垩白米663粒,碎米334粒。

2.2 实验环境和参数设置

实验均由单计算机完成,计算机配置GPU型号为GeForce GTX 1650 Ti,16G运行内存,操作系统为64位Windows10,深度学习框架为Pytorch 1.10.1、Torchvision 0.11.2,Python3.8编程语言。

训练过程中的批量尺寸设置为8,图片输入尺寸为640×640,总训练轮数设置为200,多线程设置为3,学习速率值为0.001,动量参数值为0.9,权重衰减系数值为0.000 5,置信度阈值为0.6。

2.3 评价指标

对大米分类识别检测任务采用准确率P(Precision)、召回率R(Recall)、F1值(F1 score)、平均精度和平均精度均值作为模型的评价指标。准确率P、召回率R、F1值、平均精度和平均精度均值计算如式(4)～式(8)所示。

(4)

(5)

(6)

(7)

(8)

式中:TP为被正确划分到正样本的数量;FP 为被错误划分到正样本的数量;FN为被错误划分到负样本的数量;N为类别总数。

AP衡量模型在每个类别上的识别精度,mAP衡量模型在所有类别上的平均识别精度[19]。实验中用IOU阈值为0.5时的AP和mAP来衡量模型效果[20]。检测时间以网络检测单张图片所消耗的平均时间为标准。

3 实验结果

3.1 CG-YOLOv5网络的检测结果

由于改进网络CG-YOLOv5是基于YOLOv5s网络,同时添加小目标检测层和CBAM注意力机制和Ghost模块进行改进的,为证明其有效性,对改进前后目标检测网络进行实验对比分析。分别在光照较暗、光照较强、密集情况等室内环境下,利用YOLOv5s和CG-YOLOv5对大米图像进行识别,结果如图5所示。在光照较暗和光照较强情况下,YOLOv5s和CG-YOLOv5网络能对大米实现正确的识别,但当大米密集情况下,YOLOv5s网络误将2粒米识别成1粒,而CG-YOLOv5能准确地识别出来。检测大米目标的各类型准确率、召回率、F1值、平均精度、平均精度均值如表2所示。

表2 改进模型前后对比实验结果

图5 改进模型前后识别效果图

由表2可知,在检测时间相同的情况下研究提出的CG-YOLOv5目标检测网络在碎米的AP上比改进前的YOLOv5s网络高3.7%,在准确率上提高9.7%,这说明小目标检测层对碎米检测有很好的效果。CG-YOLOv5相较于YOLOv5s网络在平均F1精度提升2.6个百分点,mAP提升4.3个百分点,所有性能提升比较明显。CG-YOLOv5在3类大米的检测中对碎米的检测效果较差,原因是碎米面积小,在图片中所占像素点较少,因此算法对碎米检测效果不佳。CG-YOLOv5在3类大米的检测中对正常米的检测效果最好,原因是正常米特征单一且在图片所占像素点较多,因此网络对正常米有很好的检测效果。

3.2 不同注意力机制的性能对比

在基于CG-YOLOv5网络的基础上,采用了不同注意力机制进行实验。将训练完成后的模型对大米图像测试集进行验证,结果如表3所示。

表3 不同注意力机制的检测结果

从表3可以看出,有CBAM注意力机制的模型平均F1值和mAP都是最高的,性能要优于CA、SE、SAM注意力机制的模型,相较于CA注意力机制模型平均F1值和mAP分别提升3%,所以本文选择添加CBAM注意力机制。在平均F1值和mAP上,CBAM注意力机制比SAM分别提高2.8%和2%,这得益于通道注意力模块。

3.3 特征可视化

为了直观CBAM注意力机制的效果,本研究对27层结构(如表1所示)的特征图进行可视化,如图6所示,使用CBAM注意力机制后获取的特征图中大米目标轮廓信息更明显,说明CBAM增强了对大米目标信息的提取,忽略其他无关特征。

图6 测试图像和加入CBAM机制前后特征图可视化

3.4 模型检测效果对比实验

利用相同数据集对另外8种目标检测网络进行训练,得到最优权重文件后,在测试集上进行测试,9种网络的AP、mAP、平均检测时间如表4所示。在检测精度上,该方法的mAP为96.5%,均高于其他8种目标检测网络,比Faster R-CNN的mAP高40.9个百分点,比SSD高2.8个百分点;特别是对正常米的识别精度明显优于其他网络,并且,其他的8种目标检测网络对正常米的检测精度要高于碎米的平均精度。在检测时间上,该方法的平均检测时间为41 ms,均快于其他目标检测网络,比YOLOv3[21]快119 ms,比YOLOv4-tiny[22]快35 ms。综合比较各个类别的AP、mAP和平均检测时间,本研究改进的网络具有更好的性能,能够实时检测大米,满足大米品质检测时的实时性要求。

表4 检测网络性能指标对比

3.5 消融实验

为探究不同模块对网络性能的影响,做了消融实验,性能指标如表5所示。模型的复杂度可以用参数量来衡量,参数量越小,说明模型复杂度越低。采用Ghost模块的网络,在网络参数量上明显低于未使用Ghost模块的网络,说明Ghost模块具有轻量化的作用,可有效地提高模型的检测速度。加入CBAM注意力机制的网络在检测mAP上比YOLOv5s网络高2.9%,说明CBAM注意力机制能增强网络对大米的特征提取。将检测层增加为4个检测层,虽然降低了检测效率,但有效地提高了检测精度,mAP值提高了3%。引入Ghost模块的网络有效地减少了模型的参数量和检测时间,mAP提升了1.8%,充分说明了Ghost模块不仅能减少参数量,而且能提升模型的性能。实验结果表明,本文所提出的改进对网络性能均有促进作用,CG-YOLOv5与YOLOv5s网络相比mAP提高了4.3%,正常米、垩白米、碎米的AP分别提高了7.0%、2.2%、3.7%。

表5 消融实验结果

4 结论

研究提出了一种基于改进YOLOv5s网络的大米外观品质识别模型(CG-YOLOv5),用于检测大米的外观品质识别。在YOLOv5s目标检测网络的基础上,添加了小目标检测层和使用Ghost轻量化模块、融合CBAM注意力机制,提高了网络结构识别精度。

在相同数据集上实验结果表明,改进后的YOLOv5s模型的平均F1值为92.6%,mAP为96.5%,比YOLOv5s网络分别提升了2.6%和4.3%。通过实验研究不同注意力机制对网络性能的影响,表明CBAM注意力机制对网络性能更有促进作用。为证明CBAM注意力机制的有效性,进行了特征图可视化,说明CBAM增强了对大米目标信息的提取。然后与主流的8种网络进行实验对比,本研究改进的网络平均检测精度最高,效果最好,可以用于大米品质检测中。