基于改进YOLOv5的奶牛乳头目标检测与数据集构建

夏事成, 王 磊,2*, 王成军, 席横流

(1.安徽科技学院 机械工程学院,安徽 凤阳 233100;2.巢湖学院 电子工程学院,安徽 巢湖 238024;3.安徽理工大学 人工智能学院,安徽 淮南 232001)

目标检测作为计算机视觉领域的核心问题之一[1-4],是找出图像中所有感兴趣的目标,确定它们的类别和位置,其检测效果的准确性和实时性是衡量机器视觉的重要指标。目标检测已在众多领域得到广泛的运用[5-10]。但由于在目标检测过程中存在外界因素的干扰,目标检测依然是计算机视觉领域最具有挑战性的问题。目前主流的检测方法是基于深度学习的方法。分为两种[11]:一种是存在候选区域的二阶目标检测法(Tow Stage),包括经典的R-CNN、SPP-Net、faster R-CNN等算法。另一种是不存在候选区域的一阶目标检测法(One Stage),主要包含YOLO、SSD、RetinaNet等算法[12-15]。其中,YOLO算法因其迭代迅速、检测速度快,加快了目标检测领域中实时性问题的解决,因而广受欢迎[16-17]。

随着机器视觉技术的不断发展,科研人员一直致力于YOLO算法的维护与改进。Lalak等[18]使用YOLO算法从无人机图像中自动检测非典型航空障碍物;Wang等[19]将YOLOv5与注意机制结合通过计算重叠率来构建数据集,提出YOLO_CBAM网络模型的精度与召回率均略高于传统的YOLOv5算法;Yue等[20]对YOLO算法增加了一个新的预测层YOLO_head,实现相比于YOLOv4算法对小目标检测的平均精度(AP)提高约12%,总平均精度提高约3%;Yuan等[21]基于YOLOv5增加了注意力机制突出特征图中的重要信息,弱化次要信息,增强了网络的识别能力。

YOLOv5算法的发展虽然迅速,但在环境比较复杂以及目标比较密集的场景下,因其计算复杂度较大,导致检测速度降低,检测的实时性变差[22]。针对此问题,基于YOLOv5对中大目标的检测提出CN-YOLOv5算法,并基于YOLOv5s创建奶牛乳头图像数据集CN-YOLOv5s model,对改进后的算法进行效果验证。

1 YOLOv5算法

YOLOv5算法的识别过程主要包括输入端、主干(Backbone)、脖子(Neck)以及输出端(Prediction)等4个部分,其网络模型如图1所示。

图1 YOLOv5网络模型图Fig.1 YOLOv5 network model

输入端包括Mosaic数据增强、自适应锚框计算、自适应图片缩放等3个部分,达到预处理图片、提高检测速度的目的。在日常生活中小目标的数量一般比中目标和大目标低很多。而Coco数据集包含大量的小目标,但是小目标的分布并不均匀,因此使用Mosaic数据增强处理,在提高网络鲁棒性、小目标检测准确性、降低设备检测要求、加速检测速度等方面起到了良好效果。

网络主干部分主要是Focus结构。Focus结构是对输入端输入的图片进行一个类似切面的操作(图2)。这种结构可以有效保留图片中更多的特征。再利用3×3,步长为2的卷积核对图片进行卷积以减少信息梯度重复,起到减小计算量并增强图片梯度表现的作用。

图2 切片结构示意图Fig.2 Schematic diagram of the slice structure

YOLOv5网络结构脖子部分的重点是FPN(Feature Pyramid Network)与PAN(Path Aggregation Network)结构。FPN是通过对图片特征的上采样方式将高层信息进行传递融合,得到进行预测的特征图。PAN是把特征图进行步幅为2的3×3卷积,并与上一层下采样后的层进行横向连接,二者进行特征融合,以此类推。YOLOv5网络结构的输出端主要包含CIOU_Loss损失函数与NMS(Non Maximum Suppression)非极大值抑制两个部分。CIOU损失函数如图3所示。

图3 CIOU示意图Fig.3 Schematic diagram of CIOU

(1)

(2)

图3中c是预测框和真实框的最小外接矩阵对角线距离;d是预测框与真实框的中心点距离;v是用来衡量长宽比一致性的参数。式(2)中w为宽,h为高,gt代表真实框。YOLOv5中非极大值抑制是通过将图片分割成多个网格。每个网格的大小相同,并且让多个网格每个都可以预测出多个预测框。预测出来的每个边界框都包含物体的中心位置(x,y),物体的高h,物体的宽w,预测的置信度和预测类别等5个信息。

2 YOLOv5网络结构的改进

基于YOLOv5算法,对原始模型与网络结构进行优化,以更优秀地适应奶牛乳头目标检测,具体改进方法:(1)在主干部分后增加SE注意力模块结构,该结构通过全连接网络,根据损失将主干部分获取的特征权重再进行自主学习,对特征重新校准,增大有效特征通道的权重,提高识别效率。但也增加了一些参数和计算量。(2)将网络结构颈部中的SPP结构改进为SPPF结构,SPPF利用多个小尺寸池化核连接的方法代替了SPP中等差增大的单池化,大幅度减少计算量,达到在不影响识别准确度的同时,加快奶牛乳头识别速度的结果。(3)挤奶机的相机采集到的奶牛乳头照片中几乎没有小目标,因此在原PAN结构中,轻量化对小目标检测的特征融合过程,起到提升检测速度的作用。

2.1 YOLOv5网络结构中SE注意力模块的引入

如图4所示,在主干部分后增加SE注意力模块,YOLOv5原模型使用上采样获得大尺寸的特征图,并和主干网中相同尺寸的特征图以拼接的方式实现目标信息的聚合。SE注意力模块可以建立特征通道之间的相互依赖关系并执行特征重新校准,通过自适应学习使用全局信息选择性地强调可靠的信息性特征,并抑制对当前任务不太可靠的特征。达到对目标更加准确的识别。

图4 CN-YOLOv5网络结构模型图Fig.4 CN-YOLOv5 network structure model

SE通过压缩(Squeeze)模块和激励(Exciation)模块实现所述功能。如图5所示,SE注意力模块主要组成分为Squeeze和Excitation两部分。输入特征图大小为W×H×C,C表示通道数,W、H表示特征图宽、高。

在输入端输入数据后通过Squeeze操作,对每个特征图做全局平均池化,平均成一个实数值,如图6所示。

图6 SE模块中的squeeze示意图Fig.6 Schematic diagram of squeeze in SE model

Excitaton操作如图7所示,经过Squeeze操作后,网络输出1×1×C大小的特征图,利用权重W来学习C个通道的相关性。

得到1×1×C向量后将SE模块计算出来的各通道权重值分别和原特征图对应通道的二维矩阵相乘,得出的结果输出。原有特征向量为W×H×C,SE模块的引入一方面降低了网络计算量,一方面增加了网络的非线性能力。实现提升重要特征、抑制不重要特征的功能。

2.2 SPP结构的改进

YOLO采用如图8所示的统一步长但不同尺寸的卷积核进行卷积实现SPP,统一步长代表输出的特征图尺寸一样,只是对区域的敏感性不一样,再通过Concat函数按通道拼接后用1×1卷积,实现特征融合。

如图9所示的SPPF结构用3个池化核大小分别为5×5、9×9、13×13,步长为1的最大池化串行的方法,大幅降低计算量,使SPPF结构相较于SPP结构,在不影响检测准确度的情况下,实现更快的目标检测。

2.3 CN-PAN结构的引入

在挤奶机器人工作时,奶牛处于挤奶工位上,挤奶机器人采集到的奶牛乳头照片中,几乎不会出现小目标的情况。根据这一特点,对YOLOv5的网络结构模型进行优化,简化了对小目标识别的过程。具体结构如图10所示。

图10 CN-PAN结构示意图Fig.10 Schematic diagram of CN-PAN structure

3 实验与分析

在Windows 10 操作系统中使用YOLOv5框架进行网络模型部署,所用CPU为Intel Core i7 9750H处理器,显卡为NVIDIA GeForce GTX1650;Torch版本为1.12.1;opencv-python版本为4.6.0.66。

3.1 数据集创建

数据集训练图片共208张。为保证实验数据集的可靠性,对不同姿势与位置、密集奶牛乳头及不同遮挡程度的奶牛乳头进行采集。使用Labeling软件进行手动标注,标注为Pascalvoc格式。将数据集的80%作为训练集,20%作为测试集。模型训练输入图像尺寸设置为640×640,epochs为300,batch-size为8。

3.2 评估指标

使用目标检测通过评估指标定量分析改进YOLOv5算法对奶牛乳头检测的性能,分别为精确率(Precision)、召回率(Recall)、平均精度率(MAP)、验证损失Validation-Loss、检测速度FPS。设定置信度阈值为0.45。

(3)

(4)

(5)

其中,P表示所有检测出的目标检测正确率;R表示所有正样本中检测正确率;AP表示不同召回率下精确率的均值;TP为正样本预测为正类个数;FP为负样本预测为正类个数;FN为正样本预测为负类个数。

利用YOLOv5算法建立对奶牛乳头检测的模型库CN-YOLOv5 model。模型指标如表1所示。

表1 模型评价指标

检测结果如图11所示,第一行3张为原图,第二行3张为检测后的图片。

图11 奶牛乳头检测效果Fig.11 Detection effects of cows' nipples

3.3 模型改进前后检测速度指标

选取204张图片作为样本。利用CN-YOLOv5 model数据集对比CN-YOLOv5s与YOLOv5s对这204张图片的检测速度。测试结果如表2所示,改进后的算法相较于原算法,在相同的检测环境中检测速度提升了约13%。

表2 2种算法在奶牛乳头模型的检测结果

4 结论

本研究为实现对奶牛乳头的快速识别与检测,针对挤奶机器人拍摄到的奶牛乳头大小及形状特性展开研究,提出CN-YOLOv5算法。该算法在不影响识别准确度的同时,加快了奶牛乳头识别速度。同时利用YOLOv5算法,建立牛乳头检测模型库CN-YOLOv5 model。模型指标中Precision为0.84,Recall为0.45,MAP为0.49,Box_loss为0.03,并用此模型库进行对比检测。实验结果表明,改进后的YOLOv5算法相对于传统YOLOv5算法在几乎不影响准确率的情况下,检测速度提高约13%。本研究通过对YOLOv5的改进,实现了在复杂环境和密集目标场景中对挤奶机器人拍摄到的奶牛乳头进行快速识别,并创建奶牛乳头数据集。为挤奶机器人的智能化、自动化提供有益的技术支撑。