STN 也来卷 YOLO 了,提升图像检测的质量,并可用于下游应用 !

2024-08-08 11:27:57 浏览数 (4)

目标检测在计算机视觉领域扮演着至关重要的角色,它可以自主识别并定位感兴趣的目标。You Only Look Once(YOLO)模型是一种有效的单次检测器。然而,YOLO在杂乱或部分遮挡的场景中面临挑战,并且在处理小尺寸、低对比度物体时可能会遇到困难。 作者提出了一种将空间变换网络(STNs)集成到YOLO中的新方法,以提高性能。所提出的STN-YOLO旨在通过在检测过程之前关注图像的重要区域并提高模型的空间不变性,从而增强模型的有效性。作者提出的方法在质量和数量上都提高了目标检测性能。 作者探讨了STN模块内不同定位网络的影响,以及模型在不同空间变换下的鲁棒性。作者将STN-YOLO应用于农业目标检测的基准数据集,以及一个来自最先进的植物表型温室设施的全新数据集。

1 Introduction

植物表型分析对于作物改良[1]、产量优化[2]和可持续实践[3]至关重要。人工智能(AI),特别是目标检测算法,已经改变了植物表型分析,提高了效率和性能[4]。You Only Look Once(YOLO)已经在各种农业应用中有效地使用,例如害虫检测[5]、作物疾病检测[6]和作物收获[7]。尽管有大量的使用案例,但由于各种空间变换[8],YOLO仍有一些局限性。空间变换网络[9]是一种提高人工神经网络对空间变换鲁棒性的方法。作者 Proposal 将STN与YOLO模型集成,以引入空间不变性。STN对图像应用可学习的仿射变换,这将有助于目标检测。STN-YOLO模型展示了空间不变性,并在几个农业基准数据集上超越了基本的YOLO模型。

此外,作者提供了一个新的、高质量的、标注详细的多种植物数据集,以推进农业植物检测和表型研究。

新的植物生长和表型(PGP)数据集在以下方面与其他可用的图像数据集不同:

(1)在不同高度捕获的多光谱图像;

(2)在不同光照条件下,多种作物的挑战性图像;

(3)使用Segment Anything模型[10]辅助的精确标注;

(4)植物的大小和形状的多样性。PGP数据集的特征为农业目标检测提出了新的挑战。作者对所提出的数据库和新模型(STN-YOLO)进行了大量的目标检测实验。

作者工作的主要贡献如下:

  • 在YOLO中集成STN,以提高目标检测性能
  • 为植物目标检测提供新的基准数据集。

II Related Work

Object Detection Models

图1:作者的目标检测框架利用空间变换网络(STNs)提高网络在目标检测中的空间不变性。STNs由三个组成部分构成:定位网络、网格生成器和采样器。定位网络输出仿射变换参数(),这些参数被传递给网格生成器。网格变换应用于输入图像,从而产生新的输出图像。

目标检测是计算机视觉领域的一个突出挑战,研究者们使用深度学习来提升性能[11, 12, 13]。两阶段检测器,如Faster R-CNN[14],包含不同的阶段,具有区域 Proposal 网络(RPN)和感兴趣区域(ROI)池化,用于候选边界框[11]。相比之下,单阶段检测器,如单次检测器(SSD)[15],通过使用网格框和 Anchor 点直接预测边界框。值得注意的是,YOLO成为单阶段检测器的一个突出例子。卷积神经网络(CNN)架构的进步带来了诸如 Anchor 点框等创新,以增强目标检测[13]。关键进展包括R-CNN,它将区域 Proposal 与使用支持向量机(SVM)分类和边界框回归的CNN相结合[16]。然而,R-CNN带来了相当大的计算成本和信息损失。Fast R-CNN通过ROI池化和 Proposal 细化扩展了这一方法,提高了测试速度和算法精度[17]。Faster R-CNN进一步减少了计算需求,通过引入RPN进行区域 Proposal ,基于特定标准选择 Anchor 点,以达到卓越的识别精度[14]。持续的研究进展集中在速度优化上,包括 Anchor-Free 点检测[18]。

Spatial Transformer Network (STN)

空间变换网络(STN)是一个具有可微性质的模块,它对输入特征图执行可学习的空间变换[9]。该变换依赖于特定的输入,生成单一输出特征图。在多通道输入的情况下,每个通道都应用相同的变换。STN模块由三个组成部分构成,如图1所示。第一个组成部分是定位网络,它通过几个隐藏层传递输入特征图,以生成应用于每个特征图的 learnable affine transformation 的参数。随后,预测的变换参数 用于构建一个采样网格——一组指示输入特征图应在何处采样的点,以生成变换后的输出。最后,特征图和采样网格作为采样器的输入,生成从指定网格点处从输入采样的输出图。

定位网络 定位网络()接收宽度为、高度为、个通道的输入特征图 。定位网络提取的特征随后传递到一个全连接层,该层输出应用于每个特征图的变换的参数:。的大小取决于参数化的变换类型[9]。定位网络函数可以是任何形式,如全连接网络或卷积神经网络(例如ResNet18[19]),但应包括一个最终回归层(即全连接层)以产生变换参数(例如,仿射变换有六个参数)。

参数化采样网格 要对输入特征图执行变换,每个生成的“像素”由在输入特征图的特定位置应用一个采样核确定。在此上下文中,“像素”指的是一般特征图的一个元素,不一定是图像。在一般情境中,生成的“像素”位于规则网格上的像素,创建输出特征图。这里,和分别表示网格的高度和宽度,是通道数,在输入和输出特征图中是相同的。

为了清楚说明,假设是一个2D仿射变换。在仿射情况下,点对点的变换如方程1所示:

在这种情境下,表示输出特征图中规则网格的目标坐标,而表示定义采样点的输入特征图中的源坐标。这些坐标之间的关系由仿射变换矩阵[9]确定。STN是一个自包含的模块,具有很强的适应性,可以无缝地集成到人工神经网络架构的任何部分。在CNN开始处加入STN,使网络能够学习如何动态变换输入图像以考虑空间变化。

III Method

Plant Growth and Phenotyping Dataset

植物生长和表型(PGP)数据集是一个公开可获得的高分辨率(512×512像素)多光谱植物图像集,这些图像来自一个先进的温室。PGP包含来自不同作物1137张图像:442张玉米,246张棉花和449张水稻植株。该数据集包括从不同高度拍摄的同一植物的图像。植物表型设施包括一个安装在机械臂上的相机,捕捉植物的侧面和顶部视图(本工作仅使用侧面视图图像)。多光谱成像系统(MSIS-AGRI-1-A),集成了MSC2-AGRI-1-A快照式多光谱相机和一个4通道LED照明器,用于在四个光谱带上收集图像:580纳米(绿光),660纳米(红光),730纳米(红边)和820纳米(近红外)。

图2:来自PGP数据集的示例图像。顶部行(2a-2c)是输入图像,底部行(2e-2f)是相关的标签。

在本研究中,红光、红边和绿光波段被合并,创建了伪“RGB”图像(如图2所示),以便与预训练的机器学习模型兼容。各波段被叠加,并对最终图像应用了最小-最大归一化。这些植物图像在2022年至2024年间在受控的温室环境中捕获。可以使用Roboflow [20]等框架应用数据增强技术(例如,旋转、裁剪)来扩展数据集。相同的框架也用于标签制作,这是一个两步过程,包括使用Segment Anything Model(SAM)[10]自动生成初始标签,然后对自动生成的标签进行手动精修。

Stn-Yolo Model

在提出的STN-YOLO模型中,首先将STN模块集成在YOLO模型的起始部分,以考虑输入图像中的空间变换。图像通过STN模块后,YOLO模型用于执行目标检测任务。如第二节B部分所述,STN可以选择任何 Backbone 网络用于定位网络。STN模块内的定位网络包括各种卷积层、线性层和激活函数。

STN模型中的定位 STN模型具有定位网络,该网络接收宽度为W、高度为H和C通道的输入特征图,并输出表示变换参数的。定位网络由用户选择,可以是浅层或更深层网络。比较了不同定位网络的性能,以评估定量(如准确度、精确度和召回率)和计算上,浅层或深层、预训练模型是否会带来更好的目标检测性能。浅层网络由一个带有最大池化和ReLU激活函数的卷积层组成,与深层预训练模型形成对比。在浅层网络的卷积层之后,使用了自适应平均池化层,将特征图的空间分辨率汇聚到所需大小(例如,)。

YOLO Backbone 网络 STN模块的输出随后被送入YOLO模型,其中P1、P2、P3、P4和P5通常指特征金字塔中不同的层级或阶段,这些是如图3所示的不同卷积层。通过从输入图像中提取不同尺度的特征来获得输出特征,以检测各种大小的目标,并捕捉细粒度细节和全局上下文。P1对应于分辨率最高的最细 Level ,而P5是分辨率较低但捕捉更多上下文的粗糙 Level 。YOLO模型中的 Head 通常指负责进行预测的网络部分。 Head 从多个金字塔 Level 获取特征,并产生目标类别、边界框坐标和其他相关信息。

图3:展示了STN-YOLO模型的架构。STN模块(以浅橙色和绿色显示)被添加到网络的起始部分,以考虑空间不变性。STN之后,图像被传递到YOLO模型,其中P1-P5是YOLO Backbone 网络的层,HEAD部分负责生成最终输出,CLS BBOX是进行目标检测的损失。

YOLOv8中的损失 训练期间使用了一个由两项组成的损失函数。目标函数包括两个部分:分类损失(CLS),它对预测目标类别时的错误进行惩罚;边界框(BBOX)损失,它对预测围绕目标边界框坐标时的错误进行惩罚。结合这些损失有助于训练模型同时提高目标检测的分类和定位方面。在单一类别检测中的类别不平衡指的是检测目标目标与背景之间的不一致,通常被认为是负类别。解决这种不平衡对于提高检测性能至关重要。YOLOv8中的边界框损失使用了完整交并比(CIOU)损失[21],它考虑了框重叠和长宽比差异,提升了回归性能。分布式Focal Loss(DFL)[22]通过在训练期间动态调整每个类别的损失,有效地解决类别差异,以对抗由类别不平衡引起的偏见。

第四节 实验结果与讨论

Experimental Setup

作者对每个模型进行了三次随机初始化实验,并且YOLO骨架网络在COCO数据集上进行预训练。作者使用了AdamW优化器,初始学习率为0.002。优化器配置了不同的参数组,每组对所有参数赋予特定的权重衰减值。所有实验的批处理大小设置为16。作者为每个模型使用了100个训练周期,并设置了早停(耐心值为50)。所有实验均采用Ultralytics [23] 框架,并在NVIDIA A100 GPU上执行。

在作者的研究中,为了评估所提出植物检测模型的有效性,作者使用了各种基准数据集。GlobalWheat2020数据集包含4,000张图像,其中3,000张用于训练,1,000张用于测试,专注于小麦检测。PlantDoc数据集由2,569张图像组成,其中2,330张用于训练,239张用于测试,涉及在不同环境下进行植物识别。MelonFlower数据集包括288张图像,其中193张用于训练,95张用于测试,专门针对甜瓜花朵的识别。这些数据集共同提供了一个健壮的评估框架,用于评估模型在多种植物种类和环境下的泛化能力。

Ablation Studies using PGP Dataset

定位网络的影响 在表1中,作者概述了各种定位网络的性能指标,强调了浅层定位网络和深层预训练模型所取得的可比较的定量结果。浅层网络和深层网络在卷积层后都使用了全局平均池化层。基于这些发现,作者选择浅层网络作为STN-YOLO模型中的定位网络。这一选择源于观察到的浅层网络在提取适合学习仿射变换的特征方面的有效性,与预训练模型的能力相似(基于类似的定量衡量标准)。浅层网络和预训练网络都与 Baseline YOLO模型进行了微调。深层网络(VGG16和ResNet18)的可学习参数数量(分别约为8200万和4900万)比浅层网络(约为500万)更多,但性能提升有限,因此在后续实验中选择了浅层网络。

空间信息的影响 为了评估在STN-YOLO模型中保留更多空间信息的影响,作者对浅层定位网络的最后一层进行了修改。具体来说,表2中评估了自适应平均池化层在不同空间尺寸、和下的效果。此调查旨在评估模型对保留空间信息变化的敏感性,并确定在提高目标检测任务性能方面的最优配置。作者考虑了不保留任何空间信息()和保留一些空间信息(、)的情况。结果表明,保留最多空间信息()可以得到最佳的目标检测性能。如果特征图包含更多的空间信息,那么模型将能够利用更多的上下文来最大化性能。然而,随着空间分辨率的增加,计算成本也会随之增加,因为参数数量将增多。

使用可解释AI进行模型分析 如表3所示,所提出的模型在平均减少误报检测方面表现得更好,因为STN-YOLO模型的精确度优于 Baseline YOLO模型。STN-YOLO模型还显示出召回率和平均精度(mAP)的改善,这表明它是一个更好的目标检测模型,如图(4b-4c)所示,与提出的STN-YOLO模型相比, Baseline YOLO模型未能消除右下角的案例区域。除了表3中的这些指标外,作者还提供了来自特征类激活图(EigenCAMs)[24]的见解,如图(4d-4e)所示。这些EigenCAMs为模型解释提供了宝贵的可视化,因为它们突出了对模型预测有重要贡献的图像的显著方面。这些结果支持了这一假设:将空间不变性融入到输入图像中可以在质量和数量上都导致更好的目标检测,因为STN-YOLO模型更加强调植物区域。

增强测试 所进行的最重要的实验之一是在未见过的增强数据上测试模型的性能,如表格III所示。为此,训练数据集中没有添加数据增强,但在测试数据集中执行了数据增强,以评估STN-YOLO模型在各种变换下的鲁棒性。在测试数据集上使用的增强方法包括随机裁剪(15%缩放)、剪切(水平和垂直)以及旋转(在之间),以模拟设施中可能出现的图像情况。作者观察到,与基准YOLO模型相比,提出的STN-YOLO模型在旋转和剪切情况下对PGP数据集给出了更好的定量结果。这个结果直观合理,因为旋转和剪切是可以被STN-YOLO模型学习的仿射变换。然而,在裁剪情况下,性能与基准YOLO方法相当,因为裁剪不是仿射变换。此外,对于不同增强组合(例如,旋转与剪切,旋转与裁剪),作者观察到在大多数数据增强情况下(八种中的六种),STN-YOLO模型在精确度值上的平均表现优于YOLO基准,这表明通过融入额外的空间信息,STN-YOLO有效地减少了错误阳性数量。

图4:在PGP数据集(图3(a))中的一个示例图像上,YOLO和STN-YOLO的示例结果。作者在图3(b)和3(c)中展示了YOLO和STN-YOLO模型在检测上的差异。YOLO模型没有捕捉到花盆中叶子的部分区域,并且错误地将背景的一部分检测为物体。EigenCAM [24]被用来展示每个模型在目标检测时关注的图像区域。与YOLO强调植物花盆的情况相反,STN-YOLO专注于包含大部分植物的图像区域。

表格IV:在基准农业数据集上,每个模型的目标检测性能指标,平均值和标准差( /- 1)是在每次实验的100个周期内随机初始化后得出的。每个数据集的最佳平均指标以粗体显示。表格III:针对单独和联合使用的不同增强方法的目标检测指标,平均值和标准差( /- 1)是在每次实验的100个周期内随机初始化后得出的。最佳平均指标以粗体显示。

STNs的不变性属性。

Benchmark Datasets Results

全球小麦2020数据集 全球小麦2020数据集[25]是一个全面的农业数据集,包含了来自不同国家的小麦头图片。从表4的结果来看,STN-YOLO模型在大多数指标(除准确度外)上的平均值更高,但当考虑到标准差时, Baseline YOLO模型的平均值与STN-YOLO模型相当。作者也可以从图4(a)中定性地观察到这一点,YOLO在全球小麦2020数据集中的检测数量比STN-YOLO要少。对于全球小麦2020数据集,图像包括了小麦 Head 的“放大”视图。在这种情况下,空间不变性不会影响目标检测结果,这一点在表4中平均目标检测指标(除准确度外)的轻微提升中得到了体现。

植物文档数据集 植物文档数据集[26]源自植物村庄数据集(PVD)[27]。该数据集是利用自动化系统,借助GoogleNet[28]和AlexNet[29]创建的。从表4中还可以观察到一些有趣的观察结果,与全球小麦2020数据集不同,在STN-YOLO模型与 Baseline 模型相比,所有目标检测指标都有所提高。特别是,STN-YOLO模型的精确度高于 Baseline 模型,这表明模型减少了检测到的假阳性数量,正如PGP数据集所注意到的。从定性的角度来看,在图4(b)中,YOLO正确地检测到了叶子,但错误地将叶子分类为番茄早疫病叶。然而,STN-YOLO检测到了叶子,并将其正确分类为番茄链格孢病。在植物文档数据集中,重点仍然是各种植物和放大图像。然而,由于缺乏预先转换或旋转的输入图像,空间不变性仍然不重要。尽管如此,空间不变性可能潜在地为此数据集提供一些好处,特别是与全球小麦2020数据集相比,某些图像提供了更“缩小”的植物视角。

甜瓜花数据集 甜瓜花数据集[30]是一个在Roboflow上可用的数据集,图像是在温室中捕获的(与PGP数据集相似)。这个数据集用于展示模型在检测小型目标(例如,花朵)和植物的不同观察几何形态方面的能力。如表4所示,与 Baseline 模型相比,STN-YOLO模型在精确度值上存在显著差异。精确度指标的显著提升,展示了模型在准确检测目标方面的增强能力。从视觉上,作者可以在图4(c)中看到这种差异,其中在甜瓜花数据集的YOLO检测中,模型以IOU为0.50检测到甜瓜花,而STN-YOLO在甜瓜花数据集的检测中,以IOU值为0.60获得了更好的改进检测。这种差异不仅强调了STN-YOLO模型的有效性,还突出了其在精确度是关键性能指标的应用中的潜力。由于甜瓜花数据集在本研究中评估的所有数据集中具有最大的空间变异性,STN-YOLO模型在精确度值和其他目标检测指标上表现出显著的性能提升,如表4所示。

V Conclusion

本研究致力于将空间转换网络(STN)与YOLO模型相结合,创建了STN-YOLO模型,以解决空间不变性挑战。该模型旨在提升植物图像检测的质量,并可用于下游应用,如表型特征提取。

结果显示,STN提高了模型的鲁棒性,并在数据集中减少了误报数量,这一点从更高的精确度得分中可以看出。

STN-YOLO模型在基准数据集上的性能有所提升,展示了在处理现实世界空间变换方面的潜力。

未来的工作包括将STN整合到其他目标检测模型中(例如,YOLO的未来版本[31]),开发新的目标函数以改进模型与STN的学习,扩展作者的PGP数据集,包含更多跨多种作物和不同图像条件(如照明)的图像,并融入近红外通道。

参考

[1].Spatial Transformer Network YOLO Model for Agricultural Object Detection.

1 人点赞