目标检测—YOLO系列(二 ) 解读论文与复现代码YOLOv1 PyTorch

精读论文

前言

从这篇开始,我们将进入YOLO的学习。YOLO是目前比较流行的目标检测算法,速度快且结构简单,其他的目标检测算法如RCNN系列,以后有时间的话再介绍。

本文主要介绍的是YOLOV1,这是由以Joseph Redmon为首的大佬们于2015年提出的一种新的目标检测算法。它与之前的目标检测算法如R-CNN等不同之处在于,R-CNN等目标检测算法是两阶段算法, 步骤为先在图片上生成候选框,然后利用分类器对这些候选框进行逐一的判断;而YOLOv1是一阶段算法,是端到端的算法,它把目标检测问题看作回归问题,将图片输入单一的神经网络,然后就输出得到了图片的物体边界框,即boundingbox以及分类概率等信息。下面我们就开始学习吧。


下面是一些学习资料:

论文链接:[1506.02640] You Only Look Once: Unified, Real-Time Object Detection (arxiv.org)

项目地址 :YOLO: Real-Time Object Detection (pjreddie.com)

Github源码地址:mirrors / alexeyab / darknet · GitCode


目录

Abstract—摘要

一、Introduction—前言

二、Unified Detection—统一检测

2.1 Network Design—网络设计

2.2 Training—训练

2.3 Inference—推论

2.4 Limitations of YOLO—YOLO的局限性

三、Comparison to Other Detection Systems—与其他目标检测算法的比较

四、Experiments—实验

4.1 Comparison to Other RealTime Systems—与其他实时系统的比较

4.2 VOC 2007 Error Analysis—VOC 2007误差分析

4.3 Combining Fast R-CNN and YOLO—Fast R-CNN与YOLO的结合

4.4 VOC 2012 Results—VOC 2012结果

4.5 Generalizability: Person Detection in Artwork—泛化性:图像中的人物检测

五、Real-Time Detection In The Wild—自然环境下的实时检测

六、Conclusion—结论


Abstract—摘要

翻译

我们提出的YOLO是一种新的目标检测方法。以前的目标检测方法通过重新利用分类器来执行检测。与先前的方案不同,我们将目标检测看作回归问题从空间上定位边界框(bounding box)*并*预测该框的类别概率。我们使用单个神经网络**,在一次评估中直接从完整图像上预测边界框和类别概率。由于整个检测流程仅用一个网络,所以可以直接对检测性能进行端到端的优化。

我们的统一架构速度极快。我们的基本YOLO模型以45 fps(帧/秒)的速度实时处理图像。该网络的一个较小版本——Fast YOLO,以155 fps这样惊人的速度运行,同时仍然达到其他实时检测器的两倍。与最先进的(state-of-the-art,SOTA)检测系统相比,YOLO虽然产生了较多的定位误差,但它几乎不会发生把背景预测为目标这样的假阳性(False Positive)的错误。最后,YOLO能学习到泛化性很强的目标表征。当从自然图像学到的模型用于其它领域如艺术画作时,它的表现都优于包括DPM和R-CNN在内的其它检测方法。

精读

之前的方法(RCNN系列)

(1)通过region proposal产生大量的可能包含待检测物体的potential bounding box

(2)再用分类器去判断每个bounding box里是否包含有物体,以及物体所属类别的probability或者 confidence

(3)最后回归预测

YOLO的简介:

本文将检测变为一个regression problem(回归问题),YOLO 从输入的图像,仅仅经过一个神经网络,直接得到一些bounding box以及每个bounding box所属类别的概率。

因为整个的检测过程仅仅有一个网络,所以它可以直接进行end-to-end的优化。

end-to-end: 端到端,指的是输入原始数据,输出的是最后结果,原来输入端不是直接的原始数据,而是在原始数据中提取的特征。通过缩减人工预处理和后续处理,尽可能使模型从原始输入到最终输出,给模型更多可以根据数据自动调节的空间,增加模型的整体契合度。在CV中具体表现是,神经网络的输入为原始图片,神经网络的输出为(可以直接控制机器的)控制指令。


一、Introduction—前言

翻译

人们只需瞄一眼图像,立即知道图像中的物体是什么,它们在哪里以及它们如何相互作用。人类的视觉系统是快速和准确的,使得我们在无意中就能够执行复杂的任务,如驾驶。快速且准确的目标检测算法可以让计算机在没有专门传感器的情况下驾驶汽车,使辅助设备能够向人类用户传达实时的场景信息,并解锁通用、响应性的机器人系统的潜能。

目前的检测系统通过重用分类器来执行检测。为了检测目标,这些系统为该目标提供一个分类器,在测试图像的不同的位置和不同的尺度上对其进行评估。像deformable parts models(DPM,可变形部分模型)这样的系统使用滑动窗口方法,其分类器在整个图像上均匀间隔的位置上运行[10]。

最近的方法,如R-CNN使用region proposal(区域候选)策略,首先在图像中生成潜在的边界框(bounding box),然后在这些框上运行分类器。在分类之后,执行用于细化边界框的后处理,消除重复的检测,并根据场景中的其它目标为边界框重新打分[13]。这些复杂的流程是很慢,很难优化的,因为每个独立的部分都必须单独进行训练

**我们将目标检测看作是一个单一的回归问题,直接从图像像素得到边界框坐标和类别概率。**使用我们的系统——You Only Look Once(YOLO),便能得到图像上的物体是什么和物体的具体位置。

YOLO非常简单(见图1),它仅用单个卷积网络就能同时预测多个边界框和它们的类别概率。YOLO在整个图像上训练,并能直接优化检测性能。与传统的目标检测方法相比,这种统一的模型下面所列的一些优点

第一,YOLO速度非常快。由于我们将检测视为回归问题,所以我们不需要复杂的流程。测试时,我们在一张新图像上简单的运行我们的神经网络来预测检测结果。在Titan X GPU上不做批处理的情况下,YOLO的基础版本以每秒45帧的速度运行,而快速版本运行速度超过150fps。这意味着我们可以在不到25毫秒的延迟内实时处理流媒体视频。此外,YOLO实现了其它实时系统两倍以上的平均精度。关于我们的系统在网络摄像头上实时运行的演示,请参阅我们的项目网页:YOLO: Real-Time Object Detection。

第二,YOLO是在整个图像上进行推断的。与基于滑动窗口和候选框的技术不同,YOLO在训练期间和测试时都会顾及到整个图像,所以它隐式地包含了关于类的上下文信息以及它们的外观。Fast R-CNN是一种很好的检测方法[14],但由于它看不到更大的上下文,会将背景块误检为目标。与Fast R-CNN相比,YOLO的背景误检数量少了一半。

第三,YOLO能学习到****目标的泛化表征(generalizable representations of objects)。把在自然图像上进行训练的模型,用在艺术图像进行测试时,YOLO大幅优于DPM和R-CNN等顶级的检测方法。由于YOLO具有高度泛化能力,因此在应用于新领域或碰到意外的输入时不太可能出故障。

YOLO在精度上仍然落后于目前最先进的检测系统。虽然它可以快速识别图像中的目标,但它在定位某些物体尤其是小的物体上精度不高。

我们在实验中会进一步探讨精度/时间的权衡。我们所有的训练和测试代码都是开源的,而且各种预训练模型也都可以下载。


精读

之前的研究:

DPM: 系统为检测对象使用分类器,并在测试图像的不同位置和尺度对其进行评估

**R-CNN:**SS方法提取候选框+CNN+分类+回归。

YOLO处理步骤:

img

(1)将输入图像的大小调整为448×448,分割得到7*7网格;

(2)通过CNN提取特征和预测;

(3)利用非极大值抑制(NMS)进行筛选


YOLO的定义:

YOLO将目标检测重新定义为单个回归问题从图像像素直接到边界框坐标和类概率。YOLO可以在一个图像来预测:哪些对象是存在的?它们在哪里?

如 Figure 1:将图像输入单独的一个 CNN 网络,就会预测出 bounding box,以及这些 bounding box 所属类别的概率。

YOLO 用一整幅图像来训练,同时可以直接优化性能检测。

性能检测对比:

img


YOLO的优点:

(1)**YOLO的速度非常快。**能够达到实时的要求。在 Titan X 的 GPU 上 能够达到 45 帧每秒。

(2)**YOLO在做预测时使用的是全局图像。**与FastR-CNN相比,YOLO产生的背景错误数量不到一半。

(3)**YOLO 学到物体更泛化的特征表示。**因此当应用于新域或意外输入时,不太可能崩溃。


二、Unified Detection—统一检测

网格单元

翻译

我们将目标检测的独立部分(the separate components )整合到单个神经网络中。我们的网络使用整个图像的特征来预测每个边界框。它还可以同时预测一张图像中的所有类别的所有边界框。这意味着我们的网络对整张图像和图像中的所有目标进行全局推理(reason globally)。YOLO设计可实现端到端训练和实时的速度,同时保持较高的平均精度。

我们的系统将输入图像分成 S×S 的网格。如果目标的中心落入某个网格单元(grid cell)中,那么该网格单元就负责检测该目标。

每个网格单元都会预测 B个 边界框和这些框的置信度分数(confidence scores)。这些置信度分数反映了该模型对那个框内是否包含目标的置信度,以及它对自己的预测的准确度的估量。在形式上,我们将置信度定义为 confidence=Pr(Object)∗IOUpred truth 。如果该单元格中不存在目标(即Pr(Object)=0),则置信度分数应为 0 。否则(即Pr(Object)=1),我们希望置信度分数等于预测框(predict box)与真实标签框(ground truth)之间联合部分的交集(IOU)。

每个网格单元还预测了C类的条件概率,Pr(Classi|Object)。这些概率是以包含目标的网格单元为条件的。我们只预测每个网格单元的一组类别概率,而不考虑框B的数量。

在测试时,我们将条件类概率和单个框的置信度预测相乘:Pr(Class_{i}|Object)*Pr(Object)*IOU{truth}_{pred}=Pr(Class_{i})*IOU{truth}_{pred}

这给我们提供了每个框的特定类别的置信度分数。这些分数既是对该类出现在框里的概率的编码,也是对预测的框与目标的匹配程度的编码。


精读

思想

YOLO将目标检测问题作为回归问题。会将输入图像分成S×S的网格,如果一个物体的中心点落入到一个cell中,那么该cell就要负责预测该物体,一个格子只能预测一个物体,会生成两个预测框。

对于每个grid cell:

(1)预测B个边界框,每个框都有一个置信度分数(confidence score)这些框大小尺寸等等都随便,只有一个要求,就是生成框的中心点必须在grid cell里

(2)每个边界框包含5个元素:(x,y,w,h)

img

● x,y: 是指bounding box的预测框的中心坐标相较于该bounding box归属的grid cell左上角的偏移量,在0-1之间。

img

在上图中,绿色虚线框代表grid cell,绿点表示该grid cell的左上角坐标,为(0,0);红色和蓝色框代表该grid cell包含的两个bounding box,红点和蓝点表示这两个bounding box的中心坐标。有一点很重要,bounding box的中心坐标一定在该grid cell内部,因此,红点和蓝点的坐标可以归一化在0-1之间。在上图中,红点的坐标为(0.5,0.5),即x=y=0.5,蓝点的坐标为(0.9,0.9),即x=y=0.9。

● w,h: 是指该bounding box的宽和高,但也归一化到了0-1之间,表示相较于原始图像的宽和高(即448个像素)。比如该bounding box预测的框宽是44.8个像素,高也是44.8个像素,则w=0.1,h=0.1。

img

红框的x=0.8,y=0.5,w=0.1,h=0.2。

(3)不管框 B 的数量是多少,只负责预测一个目标

(4)预测 C 个条件概率类别(物体属于每一种类别的可能性)

img

**综上,S×S 个网格,每个网格要预测 B个bounding box (中间上图),还要预测 C 个类(中间下图)。**将两图合并,网络输出就是一个 S × S × (5×B+C)。(S x S个网格,每个网格都有B个预测框,每个框又有5个参数,再加上每个网格都有C个预测类)

Q1:为什么每个网格有固定的B个bounding box?(即B=2)

在训练的时候会在线地计算每个predictor预测的bounding box和ground truth的IOU,计算出来的IOU大的那个predictor,就会负责预测这个物体,另外一个则不预测。这么做有什么好处?我的理解是,这样做的话,实际上有两个predictor来一起进行预测,然后网络会在线选择预测得好的那个predictor(也就是IOU大)来进行预测。

Q2:每个网格预测的两个bounding box是怎么得到的?

YOLO中两个bounding box是人为选定的(2个不同 长宽比)的box,在训练开始时作为超参数输入bounding box的信息,随着训练次数增加,loss降低,bounding box越来越准确。Faster RCNN也是人为选定的(9个 不同长宽比和scale),YOLOv2是统计分析ground true box的特点得到的(5个)。


预测特征组成

最终的预测特征由边框的位置、边框的置信度得分以及类别概率组成,这三者的含义如下:

  • 边框位置: 对每一个边框需要预测其中心坐标及宽、高这4个量, 两个边框共计8个预测值边界框宽度w和高度h用图像宽度和高度归一化。因此 x,y,w,h 都在0和1之间。
  • 置信度得分(box confidence score) c : 框包含一个目标的可能性以及边界框的准确程度。类似于Faster RCNN 中是前景还是背景。由于有两个边框,因此会存在两个置信度预测值。
  • 类别概率: 由于PASCAL VOC数据集一共有20个物体类别,因此这里预测的是边框属于哪一个类别。
注意
  • 一个cell预测的两个边界框共用一个类别预测, 在训练时会选取与标签IoU更大的一个边框负责回归该真实物体框,在测试时会选取置信度更高的一个边框,另一个会被舍弃,因此7×7=49个gird cell最多只能预测49个物体。

  • 因为每一个 grid cell只能有一个分类,也就是他只能预测一个物体,这也是导致YOLO对小目标物体性能比较差的原因。

    如果所给图片极其密集,导致 grid cell里可能有多个物体,但是YOLO模型只能预测出来一个,那这样就会忽略在本grid cell内的其他物体。


    2.1 Network Design—网络设计

翻译

我们将此模型作为卷积神经网络来实现,并在Pascal VOC检测数据集[9]上进行评估。网络的初始卷积层从图像中提取特征,而全连接层负责预测输出概率和坐标。

我们的网络架构受图像分类模型GoogLeNet的启发[34]。我们的网络有24个卷积层,后面是2个全连接层。我们只使用1×1降维层,后面是3×3卷积层,这与Lin等人[22]类似,而不是GoogLeNet使用的Inception模块。

我们还训练了快速版本的YOLO,旨在推动快速目标检测的界限。快速YOLO使用具有较少卷积层(9层而不是24层)的神经网络,在这些层中使用较少的卷积核。除了网络规模之外,基本版YOLO和快速YOLO的所有训练和测试参数都是相同的。

我们网络的最终输出是7×7×30的预测张量。


精读

网络结构

YOLO网络结构借鉴了 GoogLeNet (经典神经网络论文超详细解读(三)——GoogLeNet InceptionV1学习笔记(翻译+精读+代码复现))。输入图像的尺寸为448×448,经过24个卷积层,2个全连接的层(FC),最后在reshape操作,输出的特征图大小为7×7×30。

img


Q:7×7×30怎么来的?

张量剖面图

img

(图片来源:YOLO v1详细解读_yolov1详解_迪菲赫尔曼的博客-CSDN博客)

  • 7×7: 一共划分成7×7的网格。
  • 30: 30包含了两个预测框的参数和Pascal VOC的类别参数:每个预测框有5个参数:x,y,w,h,confidence。另外,Pascal VOC里面还有20个类别;所以最后的30实际上是由5x2+20组成的,也就是说这一个30维的向量就是一个gird cell的信息。
  • 7×7×30: 总共是7 × 7个gird cell一共就是7 × 7 ×(2 × 5+ 20)= 7 × 7 × 30 tensor = 1470 outputs,正好对应论文。

网络详解

(1)YOLO主要是建立一个CNN网络生成预测7×7×1024 的张量 。

(2)然后使用两个全连接层执行线性回归,以进行7×7×2 边界框预测。将具有高置信度得分(大于0.25)的结果作为最终预测。

(3)在3×3的卷积后通常会接一个通道数更低1×1的卷积,这种方式既降低了计算量,同时也提升了模型的非线性能力。

(4)除了最后一层使用了线性激活函数外,其余层的激活函数为 Leaky ReLU

(5)在训练中使用了 Dropout 与数据增强的方法来防止过拟合。

(6)对于最后一个卷积层,它输出一个形状为 (7, 7, 1024) 的张量。 然后张量展开。使用2个全连接层作为一种线性回归的形式,它输出1470个参数,然后reshape为 (7, 7, 30)


2.2 Training—训练

翻译

我们在ImageNet的1000类竞赛数据集[30]上预训练我们的卷积层。对于预训练,我们使用图3中的前20个卷积层,接着是平均池化层和全连接层。我们对这个网络进行了大约一周的训练,并且在ImageNet 2012验证集上获得了单一裁剪图像88%的top-5准确率,与Caffe模型池中的GoogLeNet模型相当。我们使用Darknet框架进行所有的训练和推断[26]。

然后我们转换模型来执行检测训练。Ren等人表明,预训练网络中增加卷积层和连接层可以提高性能[29]。按照他们的方法,我们添加了四个卷积层和两个全连接层,这些层的权重都用随机值初始化。检测通常需要细粒度的视觉信息,因此我们将网络的输入分辨率从224×224改为448×448。

模型的最后一层预测类概率和边界框坐标。我们通过图像宽度和高度来规范边界框的宽度和高度,使它们落在0和1之间。我们将边界框x和y坐标参数化为特定网格单元位置的偏移量,所以它们的值被限定在在0和1之间。

模型的最后一层使用线性激活函数,而所有其它的层使用下面的Leaky-ReLU:

\phi (x)=\left{\begin{matrix} x, : : : : : : : if : x>0 & & \0.1x, : : : : : otherwise& & \end{matrix}\right.\phi (x)=\left{\begin{matrix} x, : : : : : : : if : x>0 & & \0.1x, : : : : : otherwise& & \end{matrix}\right.

我们对模型输出的平方和误差(sum-squared error)进行优化。我们选择使用平方和误差,是因为它易于优化,但是它并不完全符合最大化平均精度(average precision)的目标。它给分类误差与定位误差的权重是一样的,这点可能并不理想。另外,每个图像都有很多网格单元并没有包含任何目标,这将这些单元格的“置信度”分数推向零,通常压制了包含目标的单元格的梯度。这可能导致模型不稳定,从而导致训练在早期就发散(diverge)。

为了弥补平方和误差的缺陷,我们增加了边界框坐标预测的损失,并减少了不包含目标的框的置信度预测的损失。 我们使用两个参数λ c o o r d λ_{coord}λcoord和λ n o o b j λ_{noobj}λnoobj来实现这一点。 我们设定λ c o o r d = 5 λ_{coord}= 5λcoord=5 和 λ n o o b j = 0.5 λ_{noobj}=0.5λnoobj=0.5。

平方和误差对大框和小框的误差权衡是一样的,而我们的错误指标(error metric)应该要体现出,大框的小偏差的重要性不如小框的小偏差的重要性。为了部分解决这个问题,我们直接预测边界框宽度和高度的平方根,而不是宽度和高度。

YOLO为每个网格单元预测多个边界框。在训练时,每个目标我们只需要一个边界框预测器来负责。若某预测器的预测值与目标的实际值的IOU值最高,则这个预测器被指定为“负责”预测该目标。这导致边界框预测器的专业化。每个预测器可以更好地预测特定大小,方向角,或目标的类别,从而改善整体召回率(recall)。

在训练期间,我们优化以下多部分损失函数:

img

注意,如果目标存在于该网格单元中(前面讨论的条件类别概率),则损失函数仅惩罚(penalizes)分类错误。如果预测器“负责”实际边界框(即该网格单元中具有最高IOU的预测器),则它也仅惩罚边界框坐标错误。

我们用Pascal VOC 2007和2012的训练集和验证数据集进行了大约 135个epoch 的网络训练。因为我们仅在Pascal VOC 2012上进行测试,所以我们的训练集里包含了Pascal VOC 2007的测试数据。在整个训练过程中,我们使用:batch size=64,momentum=0.9,decay=0.0005

我们的学习率(learning rate)计划如下:在第一个epoch中,我们将学习率从1 0 − 3 10^{-3}10−3慢慢地提高到 1 0 − 2 10^{-2}10−2。如果从大的学习率开始训练,我们的模型通常会由于不稳定的梯度而发散(diverge)。我们继续以 1 0 − 2 10^{-2}10−2 进行75个周期的训练,然后以 1 0 − 3 10^{-3}10−3 进行30个周期的训练,最后以 1 0 − 4 10^{-4}10−4 进行30个周期的训练。

为避免过拟合,我们使用了Dropout和大量的数据增强。 在第一个连接层之后的dropout层的丢弃率设置为0.5,以防止层之间的相互适应[18]。 对于数据增强(data augmentation),我们引入高达20%的原始图像大小的随机缩放和平移(random scaling and translations )。我们还在 HSV 色彩空间中以高达 1.5 的因子随机调整图像的曝光度和饱和度。


精读

预训练分类网络

在 ImageNet 1000数据集上预训练一个分类网络,这个网络使用Figure3中的前20个卷积层,然后是一个平均池化层和一个全连接层。(此时网络输入是224×224)。

Q:主干结构的输入要求必须是448x448的固定尺寸,为什么在预训练阶段可以输入224x224的图像呢?

主要原因是加入了平均池化层,这样不论输入尺寸是多少,在和最后的全连接层连接时都可以保证相同的神经元数目。

img


训练检测网络

经过上一步的预训练,就已经把主干网络的前20个卷积层给训练好了,前20层的参数已经学到了图片的特征。接下来的步骤本质就是迁移学习,在训练好的前20层卷积层后加上4层卷积层和2层全连接层,然后在目标检测的任务上进行迁移学习。

在整个网络(24+2)的训练过程中,除最后一层采用ReLU函数外,其他层均采用leaky ReLU激活函数。leaky ReLU相对于ReLU函数可以解决在输入为负值时的零梯度问题。YOLOv1中采用的leaky ReLU函数的表达式为:

\phi (x)=\left{\begin{matrix} x, : : : : : : : if : x>0 & & \0.1x, : : : : : otherwise& & \end{matrix}\right.


NMS非极大值抑制

**概念:**NMS算法主要解决的是一个目标被多次检测的问题,意义主要在于在一个区域里交叠的很多框选一个最优的。

YOLO中具体操作

(1)对于上述的98列数据,先看某一个类别,也就是只看98列的这一行所有数据,先拿出最大值概率的那个框,剩下的每一个都与它做比较,如果两者的IoU大于某个阈值,则认为这俩框重复识别了同一个物体,就将其中低概率的重置成0。

(2)最大的那个框和其他的框比完之后,再从剩下的框找最大的,继续和其他的比,依次类推对所有类别进行操作。 注意,这里不能直接选择最大的,因为有可能图中有多个该类别的物体,所以IoU如果小于某个阈值,则会被保留。

(3)最后得到一个稀疏矩阵,因为里面有很多地方都被重置成0,拿出来不是0的地方拿出来概率和类别,就得到最后的目标检测结果了。

注意: NMS只发生在预测阶段,训练阶段是不能用NMS的,因为在训练阶段不管这个框是否用于预测物体的,他都和损失函数相关,不能随便重置成0。


损失函数

损失函数包括:

localization loss -> 坐标损失

confidence loss -> 置信度损失

classification loss -> 分类损失

img

损失函数详解:

(1)坐标损失

img

  • 第一行: 负责检测物体的框中心点(x, y)定位误差。
  • 第二行: 负责检测物体的框的高宽(w,h)定位误差,这个根号的作用就是为了修正对大小框一视同仁的缺点,削弱大框的误差。

Q:为啥加根号?

img

在上图中,大框和小框的bounding box和ground truth都是差了一点,但对于实际预测来讲,大框(大目标)差的这一点也许没啥事儿,而小框(小目标)差的这一点可能就会导致bounding box的方框和目标差了很远。而如果还是使用第一项那样直接算平方和误差,就相当于把大框和小框一视同仁了,这样显然不合理。而如果使用开根号处理,就会一定程度上改善这一问题 。

img

这样一来,同样是差一点,小框产生的误差会更大,即对小框惩罚的更严重。

(2)置信度损失

img

  • 第一行: 负责检测物体的那个框的置信度误差。
  • 第二行: 不负责检测物体的那个框的置信度误差。

(3)分类损失

img 负责检测物体的grid cell分类的误差。


特殊符号的含义:

img


2.3 Inference—推论

翻译

就像在训练中一样,预测测试图像的检测只需要一次网络评估。在Pascal VOC上,每张图像上网络预测 98 个边界框和每个框的类别概率。YOLO在测试时非常快,因为它只需要一次网络评估(network evaluation),这与基于分类器的方法不同。

网格设计强化了边界框预测中的空间多样性。通常一个目标落在哪一个网格单元中是很明显的,而网络只能为每个目标预测一个边界框。然而,一些大的目标或接近多个网格单元的边界的目标能被多个网格单元定位。非极大值抑制(Non-maximal suppression,NMS)可以用来修正这些多重检测。非最大抑制对于YOLO的性能的影响不像对于R-CNN或DPM那样重要,但也能增加2−3%的mAP。


精读

(1)预测测试图像的检测只需要一个网络评估。

(2)测试时间快

(3)当图像中的物体较大,或者处于 grid cells 边界的物体,可能在多个 cells 中被定位出来。

(4)利用NMS去除重复检测的物体,使mAP提高,但和RCNN等相比不算大。


2.4 Limitations of YOLO—YOLO的局限性

翻译

由于每个格网单元只能预测两个框,并且只能有一个类,因此YOLO对边界框预测施加了很强的空间约束。这个空间约束限制了我们的模型可以预测的邻近目标的数量。我们的模型难以预测群组中出现的小物体(比如鸟群)。

由于我们的模型学习是从数据中预测边界框,因此它很难泛化到新的、不常见的长宽比或配置的目标。我们的模型也使用相对较粗糙的特征来预测边界框,因为输入图像在我们的架构中历经了多个下采样层(downsampling layers)。

最后,我们的训练基于一个逼近检测性能的损失函数,这个损失函数无差别地处理小边界框与大边界框的误差。大边界框的小误差通常是无关要紧的,但小边界框的小误差对IOU的影响要大得多。我们的主要错误来自于不正确的定位


精读

(1)对于图片中一些群体性小目标检测效果比较差。因为yolov1网络到后面感受野较大,小目标的特征无法再后面7×7的grid中体现,针对这一点,yolov2已作了一定的修改,加入前层(感受野较小)的特征进行融合。

(2)原始图片只划分为7x7的网格,当两个物体靠的很近时(挨在一起且中点都落在同一个格子上的情况),效果比较差。因为yolov1的模型决定了一个grid只能预测出一个物体,所以就会丢失目标,针对这一点,yolov2引入了anchor的概念,一个grid有多少个anchor理论上就可以预测多少个目标。

(3)每个网格只对应两个bounding box,当物体的长宽比不常见(也就是训练数据集覆盖不到时),效果较差。

(4)最终每个网格只对应一个类别,容易出现漏检(物体没有被识别到)。


三、Comparison to Other Detection Systems—与其他目标检测算法的比较

翻译

目标检测是计算机视觉中的核心问题。检测流程通常是首先从输入图像上提取一组鲁棒特征(Haar [25],SIFT [23],HOG [4],卷积特征[6])。然后,分类器[36,21,13,10]或定位器[1,32]被用来识别特征空间中的目标。这些分类器或定位器或在整个图像上或在图像中的一些子区域上以滑动窗口的方式运行[35,15,39]。我们将YOLO检测系统与几种顶级检测框架进行比较,突出了关键的相似性和差异性。

Deformable parts models。可变形部分模型(DPM)使用滑动窗口方法进行目标检测[10]。DPM使用不相交的流程来提取静态特征,对区域进行分类,预测高评分区域的边界框等。我们的系统用单个卷积神经网络替换所有这些不同的部分。网络同时进行特征提取,边界框预测,非极大值抑制和上下文推理。网络的特征feature是在**在线(in-line)**训练出来的而不是静态,因此可以根据特定的检测任务进行优化。我们的统一架构比DPM更快,更准确。

R-CNN。R-CNN及其变体(variants)使用区域候选而不是滑动窗口来查找图像中的目标。选择性搜索[35]生成潜在的边界框(Selective Search generates potential bounding boxes),卷积网络提取特征,SVM对框进行评分,线性模型调整边界框,非最大抑制消除重复检测(eliminates duplicate detections)。 这个复杂流水线的每个阶段都必须独立地进行精确调整(precisely tuned independently),所得到的系统非常缓慢,在测试时间每个图像需要超过40秒[14]。

YOLO与R-CNN有一些相似之处。每个网格单元提出潜在的边界框并使用卷积特征对这些框进行评分。然而,我们的系统对网格单元的候选框施加空间限制,这有助于缓解对同一目标的多次检测的问题。 我们的系统还生成了更少的边界框,每张图像只有98个,而选择性搜索则有约2000个。最后,我们的系统将这些单独的组件(individual components)组合成一个单一的、共同优化的模型。

其它快速检测器。 Fast R-CNN 和 Faster R-CNN 通过共享计算和使用神经网络替代选择性搜索[14],[28]来提出候选区域来加速 R-CNN 框架。虽然它们提供了比 R-CNN 更快的速度和更高的准确度,但仍然不能达到实时性能。

许多研究工作集中在加快DPM流程上[31] [38] [5]。它们加速HOG计算,使用级联(cascades),并将计算推动到(多个)GPU上。但是,实际上只有30Hz的DPM [31]可以实时运行。

YOLO并没有试图优化大型检测流程的单个组件,相反,而是完全抛弃(throws out…entirely)了大型检测流程,并通过设计来提高速度。

像人脸或行人等单个类别的检测器可以高度优化,因为他们只需处理较少的多样性[37]。YOLO是一种通用的检测器,它可以同时(simultaneously)检测多个目标。

Deep MultiBox。与R-CNN不同,Szegedy等人 训练一个卷积神经网络来预测感兴趣的区域(regions of interest,ROI)[8],而不是使用选择性搜索。 MultiBox还可以通过用单个类别预测替换置信度预测来执行单个目标检测。 但是,MultiBox无法执行一般的目标检测,并且仍然只是较大检测流水线中的一部分,需要进一步的图像补丁分类。 YOLO和MultiBox都使用卷积网络来预测图像中的边界框,但YOLO是一个完整的检测系统

OverFeat。Sermanet等人训练了一个卷积神经网络来执行定位,并使该定位器进行检测[32]。OverFeat高效地执行滑动窗口检测,但它仍然是一个不相交的系统(disjoint system)。OverFeat优化了定位功能,而不是检测性能。像DPM一样,定位器在进行预测时只能看到局部信息。OverFeat无法推断全局上下文,因此需要大量的后处理来产生连贯的检测。

MultiGrasp。我们的系统在设计上类似于Redmon等[27]的抓取检测。 我们的网格边界框预测方法基于MultiGrasp系统进行回归分析。 然而,抓取检测比物体检测要简单得多。 MultiGrasp只需要为包含一个目标的图像预测一个可抓取区域。 它不必估计目标的大小,位置或边界或预测它的类别,只需找到适合抓取的区域就可以了。 而YOLO则是预测图像中多个类的多个目标的边界框和类概率


精读

DPM

用传统的HOG特征方法,也用的是传统的支持向量机SVM分类器,然后人工造一个模板,再用滑动窗口方法不断的暴力搜索整个待识别图,去套那个模板。这个方法比较大的问题就是在于设计模板,计算量巨大,而且是个静态的,没办法匹配很多变化的东西,鲁棒性差。

R-CNN
  • 第一阶段:每个图片使用选择性搜索SS方法提取2000个候选框。
  • 第二阶段:将每个候选框送入CNN网络进行分类(使用的SVM)。

YOLO对比他们俩都很强,YOLO和R-CNN也有相似的地方,比如也是提取候选框,YOLO的候选框就是上面说过的那98个 bounding boxes,也是用到了NMS非极大值抑制,也用到了CNN提取特征。

Other Fast Detectors

Fast和Faster R-CNN :这俩模型都是基于R-CNN的改版,速度和精度都提升了很多,但是也没办法做到实时监测,也就是说FPS到不了30,作者在这里并没有谈准确度的问题,实际上YOLO的准确度在这里是不占优势的,甚至于比他们低。

Deep MultiBox

训练卷积神经网络来预测感兴趣区域,而不是使用选择性搜索。多盒也可以用单个类预测替换置信预测来执行单个目标检测。YOLO和MultiBox都使用卷积网络来预测图像中的边界框,但YOLO是一个完整的检测系统。

OverFeat

OverFeat有效地执行滑动窗口检测,优化了定位,而不是检测性能。与DPM一样,定位器在进行预测时只看到本地信息。OverFeat不能推理全局环境。

MultiGrasp

YOLO在设计上与Redmon等人的抓取检测工作相似。边界盒预测的网格方法是基于多重抓取系统的回归到抓取。

总之,作者就是给前人的工作都数落一遍,凸显自己模型的厉害(学到了!)


四、Experiments—实验

4.1 Comparison to Other RealTime Systems—与其他实时系统的比较

翻译

目标检测方面的许多研究工作都集中在使标准的检测流程更快[5],[38],[31],[14],[17],[28]。然而,只有Sadeghi等人实际上产生了一个实时运行的检测系统(每秒30帧或更好)[31]。我们将YOLO与DPM的GPU实现进行了比较,其在30Hz或100Hz下运行。虽然其它的算法没有达到实时性的标准,我们也比较了它们的mAP和速度的关系,从而探讨目标检测系统中精度和性能之间的权衡

Fast YOLO是PASCAL上最快的目标检测方法;据我们所知,它是现有的最快的目标检测器。具有52.7%的mAP,实时检测的精度是以前的方法的两倍以上。普通版YOLO将mAP推到63.4%的同时保持了实时性能。

我们还使用VGG-16训练YOLO。 这个模型比普通版YOLO更精确,但也更慢。 它的作用是与依赖于VGG-16的其他检测系统进行比较,但由于它比实时更慢,所以本文的其他部分将重点放在我们更快的模型上。

最快的DPM可以在不牺牲太多mAP的情况下有效加速DPM,但仍然会将实时性能降低2倍[38]。与神经网络方法相比,DPM的检测精度相对较低,这也是限制它的原因。

减去R的R-CNN用静态侯选边界框取代选择性搜索[20]。虽然速度比R-CNN更快,但它仍然无法实时,并且由于该方法无法找到好的边界框,准确性受到了严重影响。

Fast R-CNN加快了R-CNN的分类阶段,但它仍然依赖于选择性搜索,每个图像需要大约2秒才能生成边界候选框。因此,它虽然具有较高的mAP,但的速度是0.5 fps,仍然远未达到实时。

最近的Faster R-CNN用神经网络替代了选择性搜索来候选边界框,类似于Szegedy等人[8]的方法。在我们的测试中,他们最准确的模型达到了 7fps,而较小的、不太准确的模型以18 fps运行。 Faster R-CNN的VGG-16版本比YOLO高出10mAP,但比YOLO慢了6倍。 Zeiler-Fergus 版本的Faster R-CNN只比YOLO慢2.5倍,但也不如YOLO准确。


精读

Table 1 在Pascal VOC 2007 上与其他检测方法的对比img

**结论:**实时目标检测(FPS>30),YOLO最准,Fast YOLO最快。


4.2 VOC 2007 Error Analysis—VOC 2007误差分析

翻译

为了进一步研究YOLO和最先进的检测器之间的差异,我们详细分析了VOC 2007的分类(breakdown)结果。我们将YOLO与Fast R-CNN进行比较,因为Fast R-CNN是PASCAL上性能最高的检测器之一并且它的检测代码是可公开得到的。

我们使用Hoiem等人的方法和工具[19],对于测试的每个类别,我们查看该类别的前N个预测。每个预测都或是正确的,或是根据错误的类型进行分类:

  • Correct: correct class and IOU>0.5
  • Localization: correct class, 0.1
  • Similar: class is similar, IOU>0.1
  • Other: class is wrong, IOU>0.1
  • Background: IOU<0.1 for any object(所有目标的IOU都<0.1)

YOLO难以正确地定位目标,因此定位错误比YOLO的所有其他错误总和都要多。Fast R-CNN定位错误更少,但把背景误认成目标的错误比较多。它的最高检测结果中有13.6%是不包含任何目标的误报(false positive,背景)。 Fast R-CNN把背景误认成目标的概率比YOLO高出3倍。


精读

本文使用HoeMm等人的方法和工具。对于测试时间的每个类别,查看该类别的N个预测。每个预测要么是正确的,要么是基于错误类型进行分类的:

img

参数含义:

  • **Correct:**正确分类,且预测框与ground truth的IOU大于0.5,既预测对了类别,预测框的位置和大小也很合适。
  • **Localization:**正确分类,但预测框与ground truth的IOU大于0.1小于0.5,即虽然预测对了类别,但预测框的位置不是那么的严丝合缝,不过也可以接受。
  • Similar: 预测了相近的类别,且预测框与ground truth的IOU大于0.1。即预测的类别虽不正确但相近,预测框的位置还可以接受。
  • **Other:**预测类别错误,预测框与ground truth的IOU大于0.1。即预测的类别不正确,但预测框还勉强把目标给框住了。
  • **Background:**预测框与ground truth的IOU小于0.1,即该预测框的位置为背景,没有目标。

Figure 4 显示了所有20个类中每种错误类型的平均细分情况img

**结论:**YOLO定位错误率高于Fast R-CNN;Fast R-CNN背景预测错误率高于YOLO


4.3 Combining Fast R-CNN and YOLO—Fast R-CNN与YOLO的结合

翻译

YOLO误认背景为目标的情况比Fast R-CNN少得多。 通过使用YOLO消除Fast R-CNN的背景检测,我们获得了显著的性能提升。 对于R-CNN预测的每个边界框,我们检查YOLO是否预测了一个相似的框。 如果确实如此,那么我们会根据YOLO预测的概率和两个框之间的重叠情况提高预测值。

最好的Fast R-CNN模型在VOC 2007测试集中达到了 71.8% 的mAP。 当与YOLO合并时,其mAP增加了 3.2% 至 75.0%。 我们还尝试将顶级Fast R-CNN模型与其他几个版本的Fast R-CNN结合起来。 这写的结合的平均增长率在 0.3% 至 0.6% 之间。

结合YOLO后获得的性能提高不仅仅是模型集成的副产品,因为结合不同版本的Fast R-CNN几乎没有什么益处。 相反,正是因为YOLO在测试时出现了各种各样的错误,所以它在提高Fast R-CNN的性能方面非常有效。

不幸的是,这种组合不会从YOLO的速度中受益,因为我们分别运行每个模型,然后合并结果。 但是,由于YOLO速度如此之快,与Fast R-CNN相比,它不会增加任何显著的计算时间。

精读

Table2 模型组合在VOC 2007上的实验结果对比img

**结论:**因为YOLO在测试时犯了各种错误,所以它在提高快速R-CNN的性能方面非常有效。但是这种组合并不受益于YOLO的速度,由于YOLO很快,和Fast R-CNN相比,它不增加任何有意义的计算时间。


4.4 VOC 2012 Results—VOC 2012结果

翻译

在VOC 2012测试集中,YOLO的mAp得分是57.9%。这比现有最先进的技术水平低,更接近使用VGG-16的原始的R-CNN,见表3。与其最接近的竞争对手相比,我们的系统很难处理小物体上(struggles with small objects)。在瓶子、羊、电视/监视器等类别上,YOLO得分比R-CNN和Feature Edit低8-10%。然而,在其他类别,如猫和火车YOLO取得了更好的表现。

我们的Fast R-CNN + YOLO模型组合是性能最高的检测方法之一。 Fast R-CNN与YOLO的组合提高了2.3%,在公共排行榜上提升了5个位置。

精读

Table 3 在VOC2012上mAP排序img

**结论:**Fast R-CNN从与YOLO的组合中得到2.3%的改进,在公共排行榜上提升了5个百分点。


4.5 Generalizability: Person Detection in Artwork—泛化性:图像中的人物检测

翻译

用于目标检测的学术数据集的训练和测试数据是服从同一分布的。但在现实世界的应用中,很难预测所有可能的用例,他的测试数据可能与系统已经看到的不同[3]。我们将YOLO与其他检测系统在毕加索(Picasso)数据集[12]和人物艺术(People-Art)数据集[3]上进行了比较,这两个数据集用于测试艺术品上的人物检测。

作为参考(for reference),我们提供了VOC 2007的人形检测的AP,其中所有模型仅在VOC 2007数据上训练。在Picasso数据集上测试的模型在是在VOC 2012上训练,而People-Art数据集上的模型则在VOC 2010上训练。

R-CNN在VOC 2007上有很高的AP值。然而,当应用于艺术图像时,R-CNN显着下降。R-CNN使用选择性搜索来调整自然图像的候选边界框。R-CNN在分类器阶段只能看到小区域,而且需要有很好的候选框。

DPM在应用于艺术图像时可以很好地保持其AP。之前的研究认为DPM表现良好,因为它具有强大的物体形状和布局空间模型。虽然DPM不会像R-CNN那样退化,但它的AP本来就很低。

YOLO在VOC 2007上表现出色,其应用于艺术图像时其AP降低程度低于其他方法。与DPM一样,YOLO模拟目标的大小和形状,以及目标之间的关系和目标通常出现的位置之间的关系。艺术图像和自然图像在像素级别上有很大不同,但它们在物体的大小和形状方面相似,因此YOLO仍然可以预测好的边界框和检测结果。

精读

Figure 5 通用性(Picasso 数据集和 People-Art数据集)img

**结论:**YOLO都具有很好的检测结果


五、Real-Time Detection In The Wild—自然环境下的实时检测

翻译

YOLO是一款快速,精确的物体检测器,非常适合计算机视觉应用。 我们将YOLO连接到网络摄像头,并验证它是否保持实时性能,包括从摄像头获取图像并显示检测结果的时间。

由此产生的系统是互动的和参与的。 虽然YOLO单独处理图像,但当连接到网络摄像头时,它的功能类似于跟踪系统,可在目标移动并在外观上发生变化时检测目标。 系统演示和源代码可在我们的项目网站上找到:YOLO: Real-Time Object Detection。

精读

img

**结论:**将YOLO连接到一个网络摄像头上,并验证它是否保持了实时性能,包括从摄像头中获取图像和显示检测结果的时间。结果证明效果很好,如上图所示,除了第二行第二个将人误判为飞机以外,别的没问题。


六、Conclusion—结论

翻译

我们介绍YOLO——一种用于物体检测的统一模型。 我们的模型构造简单,可以直接在完整图像上训练。 与基于分类器的方法不同,YOLO是通过与检测性能直接对应的损失函数进行训练的,并且整个模型是一起训练的。

快速YOLO是文献中最快的通用目标检测器,YOLO推动实时对象检测的最新技术。 YOLO还能很好地推广到新领域,使其成为快速,鲁棒性强的应用的理想选择。

精读

到底什么是YOLO?
  • YOLO眼里目标检测是一个回归问题
  • 一次性喂入图片,然后给出bbox和分类概率
  • 简单来说,只看一次就知道图中物体的类别和位置
YOLO过程总结:

训练阶段:

首先将一张图像分成 S × S个 gird cell,然后将它一股脑送入CNN,生成S × S × (B × 5 + C)个结果,最后根据结果求Loss并反向传播梯度下降。

预测、验证阶段:

首先将一张图像分成 S × S网格(gird cell),然后将它一股脑送入CNN,生成S × S × (B × 5 + C)个结果,最后用NMS选出合适的预选框。


复现代码

复现YOLO v1 PyTorch

Paper: [1506.02640] You Only Look Once: Unified, Real-Time Object Detection (arxiv.org)

Github: EclipseR33/yolo_v1_pytorch (github.com)

数据集

VOC2007:The PASCAL Visual Object Classes Challenge 2007 (VOC2007)

VOC2012:The PASCAL Visual Object Classes Challenge 2012 (VOC2012)

PASCAL VOC 07/12的目录结构都是一致的,因此只需要针对VOC07编写代码再扩展即可。VOC2007目录下有5个文件夹。我们需要其中的’Annotations’(存有标注信息),‘ImageSets’(存有train、val、test各类文件名), ‘JPEGImages’(存有图像)。VOC中的图像都是.jpg文件,ImageSets中的文件都是.txt文件,Annotations中的注释都是.xml文件。

//这是一个xml注释的示例,我们需要其中的<object>信息
<annotation>
	<folder>VOC2007folder>
	<filename>000001.jpgfilename>
	<source>
		<database>The VOC2007 Databasedatabase>
		<annotation>PASCAL VOC2007annotation>
		<image>flickrimage>
		<flickrid>341012865flickrid>
	source>
	<owner>
		<flickrid>Fried Camelsflickrid>
		<name>Jinky the Fruit Batname>
	owner>
	<size>
		<width>353width>
		<height>500height>
		<depth>3depth>
	size>
	<segmented>0segmented>
	<object>
		<name>dogname>
		<pose>Leftpose>
		<truncated>1truncated>
		<difficult>0difficult>
		<bndbox>
			<xmin>48xmin>
			<ymin>240ymin>
			<xmax>195xmax>
			<ymax>371ymax>
		bndbox>
	object>
	<object>
		<name>personname>		// name中包含的就是class信息
		<pose>Leftpose>
		<truncated>1truncated>
		<difficult>0difficult>
		<bndbox>
			<xmin>8xmin>
			<ymin>12ymin>
			<xmax>352xmax>
			<ymax>498ymax>
		bndbox>
	object>
annotation>
find_classes.py

首先我们需要获得VOC数据集中的所有class信息并为其编号,将该信息存储到json文件中

# 路径: ./dataset/find_classes.py
import xml.etree.ElementTree as ET
from tqdm import tqdm
import json
import os


def xml2dict(xml):
    """
    使用递归读取xml文件
    若c指向的元素是person,那么c.tag是name,c.text则是person
    """
    # data初始化时就已经将所有子元素的tag定义为key
    data = {c.tag: None for c in xml}
    for c in xml:
        # add函数用于将tag与text添加到data中
        def add(data, tag, text):
            if data[tag] is None:
                # data中该tag为空则直接添加text
                data[tag] = text
            elif isinstance(data[tag], list):
                # data中该tag为不为空且已经创建了list则append(text)
                data[tag].append(text)
            else:
                # data中该tag不为空但是没有创建list,需要先创建list
                data[tag] = [data[tag], text]
            return data

        if len(c) == 0:
            # len(c)表示c的子元素个数,若为0则表示c是叶元素,没有子元素
            data = add(data, c.tag, c.text)
        else:
            data = add(data, c.tag, xml2dict(c))
    return data


json_path = './classes.json'	# json保存到的地址

root = r'F:\AI\Dataset\VOC2012\VOCdevkit\VOC2012'	# 数据集root(VOC2007与VOC2012的class信息一致)
# 获取所有xml注释地址
annotation_root = os.path.join(root, 'Annotations')		
annotation_list = os.listdir(annotation_root)
annotation_list = [os.path.join(annotation_root, a) for a in annotation_list]

s = set()
for annotation in tqdm(annotation_list):
    xml = ET.parse(os.path.join(annotation)).getroot()
    data = xml2dict(xml)['object']
    if isinstance(data, list):
        # 有多个object
        for d in data:
            s.add(d['name'])
    else:
        # 仅有一个object
        s.add(data['name'])

s = list(s)
s.sort()

# 以class名称为key可以便于xml2label的转换
data = {value: i for i, value in enumerate(s)}
json_str = json.dumps(data)

with open(json_path, 'w') as f:
    f.write(json_str)

运行./dataset/find_classes.py之后,我们在指定目录下得到一个json文件

{"aeroplane": 0, "bicycle": 1, "bird": 2, "boat": 3, "bottle": 4, "bus": 5, "car": 6, "cat": 7, "chair": 8, "cow": 9, "diningtable": 10, "dog": 11, "horse": 12, "motorbike": 13, "person": 14, "pottedplant": 15, "sheep": 16, "sofa": 17, "train": 18, "tvmonitor": 19}

接下来我们要开始写data.py文件,主要流程:

1.通过ImageSets中train.txt, val.txt, test.txt文件的指引寻找数据集对应的所有图像名与其地址。

2.编写getitem,读取xml中的信息并转化为label形式,读取图像,并将两者都转为tensor

这里设置的dataset传出的label都是xmin, ymin, xmax, ymax, class的VOC形式,并且是直接的坐标数值,并没有使用百分比表示。

data.py
# 路径: ./dataset/data.py
from dataset.transform import *		# 导入我们重写的transform类

from torch.utils.data import Dataset
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np
import json
import os


def get_file_name(root, layout_txt):
    with open(os.path.join(root, layout_txt)) as layout_txt:
        """
        .read()			读取文件中的数据,会得到一个str字符串
        .split('\n')	以\n回车符为分界将str字符串分割成list
        [:-1]			去除最后一个空字符串,文件末尾有\n,分割后会有空字符串所以要去除
        """
        file_name = layout_txt.read().split('\n')[:-1]
    return file_name


def xml2dict(xml):
    # 这里的xml2dict与上一个文件的xml2dict一致
    data = {c.tag: None for c in xml}
    for c in xml:
        def add(data, tag, text):
            if data[tag] is None:
                data[tag] = text
            elif isinstance(data[tag], list):
                data[tag].append(text)
            else:
                data[tag] = [data[tag], text]
            return data

        if len(c) == 0:
            data = add(data, c.tag, c.text)
        else:
            data = add(data, c.tag, xml2dict(c))
    return data


class VOC0712Dataset(Dataset):
    def __init__(self, root, class_path, transforms, mode, data_range=None, get_info=False):
        # label: xmin, ymin, xmax, ymax, class
		
        # 从json文件中获得class的信息
        with open(class_path, 'r') as f:
            json_str = f.read()
            self.classes = json.loads(json_str)
           
        """
        如果是train模式,那么root的输入将为一个list(长为2,分别为2007、2012两年的数据集根目录, main中的root0712是一个示例)。将两个root与train、val两种分割组合成四个layout_txt路径,这四个路径指向VOC07/12的所有可用训练数据。
        如果是test模式那么只有VOC2007的test分割可用。这里也转换为list形式,就可以同一两种模式的代码。
        """
        layout_txt = None
        if mode == 'train':
            root = [root[0], root[0], root[1], root[1]]
            layout_txt = [r'ImageSets\Main\train.txt', r'ImageSets\Main\val.txt',
                          r'ImageSets\Main\train.txt', r'ImageSets\Main\val.txt']
        elif mode == 'test':
            if not isinstance(root, list):
                root = [root]
            layout_txt = [r'ImageSets\Main\test.txt']
        assert layout_txt is not None, 'Unknown mode'
        

        self.transforms = transforms
        self.get_info = get_info	# get_info表示在getitem时是否需要获得图像的名称以及图像大小信息 bool
		
        # 由于有多root,所以image_list与annotation_list均存储了图像与xml文件的绝对路径
        self.image_list = []
        self.annotation_list = []
        for r, txt in zip(root, layout_txt):
            self.image_list += [os.path.join(r, 'JPEGImages', t + '.jpg') for t in get_file_name(r, txt)]
            self.annotation_list += [os.path.join(r, 'Annotations', t + '.xml') for t in get_file_name(r, txt)]
		
        # data_range是一个二元tuple,分别表示数据集需要取哪一段区间,训练时若使用全部的数据则无需传入data_range,默认None的取值是会选择所有的数据的
        if data_range is not None:
            self.image_list = self.image_list[data_range[0]: data_range[1]]
            self.annotation_list = self.annotation_list[data_range[0]: data_range[1]]

    def __len__(self):
        # 返回数据集长度
        return len(self.annotation_list)

    def __getitem__(self, idx):
        image = Image.open(self.image_list[idx])
        image_size = image.size
        label = self.label_process(self.annotation_list[idx])

        if self.transforms is not None:
            # 由于目标检测中image的变换如随机裁剪与Resize都会导致label的变化,所以需要重写transform,添加部分的label处理代码
            image, label = self.transforms(image, label)
        if self.get_info:
            return image, label, os.path.basename(self.image_list[idx]).split('.')[0], image_size
        else:
            return image, label

    def label_process(self, annotation):
        xml = ET.parse(os.path.join(annotation)).getroot()
        data = xml2dict(xml)['object']
        # 根据data的两种形式将其读取到label中,并将label转为numpy形式
        if isinstance(data, list):
            label = [[float(d['bndbox']['xmin']), float(d['bndbox']['ymin']),
                     float(d['bndbox']['xmax']), float(d['bndbox']['ymax']),
                     self.classes[d['name']]]
                     for d in data]
        else:
            label = [[float(data['bndbox']['xmin']), float(data['bndbox']['ymin']),
                     float(data['bndbox']['xmax']), float(data['bndbox']['ymax']),
                     self.classes[data['name']]]]
        label = np.array(label)
        return label


if __name__ == "__main__":
    from dataset.draw_bbox import draw

    root0712 = [r'F:\AI\Dataset\VOC2007\VOCdevkit\VOC2007', r'F:\AI\Dataset\VOC2012\VOCdevkit\VOC2012']

    transforms = Compose([
        ToTensor(),
        RandomHorizontalFlip(0.5),
        Resize(448)
    ])
    ds = VOC0712Dataset(root0712, 'classes.json', transforms, 'train', get_info=True)
    print(len(ds))
    for i, (image, label, image_name, image_size) in enumerate(ds):
        if i <= 1000:
            continue
        elif i >= 1010:
            break
        else:
            print(label.dtype)
            print(tuple(image.size()[1:]))
            draw(image, label, ds.classes)
    print('VOC2007Dataset')
transform.py
# 路径: ./dataset/transform.py
import torch
import torchvision
import random


class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, label):
        for t in self.transforms:
            image, label = t(image, label)
        return image, label


class ToTensor:
    def __init__(self):
        self.totensor = torchvision.transforms.ToTensor()

    def __call__(self, image, label):
        image = self.totensor(image)
        label = torch.tensor(label)
        return image, label


class RandomHorizontalFlip:
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, image, label):
        """
        :param label: xmin, ymin, xmax, ymax
        如果图片被水平翻转,那么label的xmin与xmax会互换,变成 xmax, ymin, xmin, ymax
        由于YOLO的输出是(center_x, center_y, w, h) ,因此label的xmin与xmax换位不会影响损失计算与训练但是需要注意w,h计算时使用abs
        """
        if random.random() < self.p:
            height, width = image.shape[-2:]
            image = image.flip(-1)      # 水平翻转
            bbox = label[:, :4]
            bbox[:, [0, 2]] = width - bbox[:, [0, 2]]
            label[:, :4] = bbox
        return image, label


class Resize:
    def __init__(self, image_size, keep_ratio=True):
        """
        :param image_size: int
        keep_ratio = True  保留宽高比
        keep_ratio = False 填充成正方形
        """
        self.image_size = image_size
        self.keep_ratio = keep_ratio

    def __call__(self, image, label):
        """
        :param in_image: tensor [3, h, w]
        :param label: xmin, ymin, xmax, ymax
        :return:
        """
        # 将所有图片左上角对齐构成448*448tensor的Transform

        h, w = tuple(image.size()[1:])
        label[:, [0, 2]] = label[:, [0, 2]] / w
        label[:, [1, 3]] = label[:, [1, 3]] / h

        if self.keep_ratio:
            r_h = min(self.image_size / h, self.image_size / w)
            r_w = r_h
        else:
            r_h = self.image_size / h
            r_w = self.image_size / w

        h, w = int(r_h * h), int(r_w * w)
        h, w = min(h, self.image_size), min(w, self.image_size)
        label[:, [0, 2]] = label[:, [0, 2]] * w
        label[:, [1, 3]] = label[:, [1, 3]] * h

        T = torchvision.transforms.Resize([h, w])

        Padding = torch.nn.ZeroPad2d((0, self.image_size - w, 0, self.image_size - h))
        image = Padding(T(image))

        assert list(image.size()) == [3, self.image_size, self.image_size]
        return image, label
draw_bbox.py
# 路径: ./dataset/draw_bbox.py
import torchvision.transforms as F
import numpy as np
from PIL import ImageDraw, ImageFont
import matplotlib.pyplot as plt


colors = ['Pink', 'Crimson', 'Magenta', 'Indigo', 'BlueViolet',
          'Blue', 'GhostWhite', 'LightSteelBlue', 'Brown', 'SkyBlue',
          'Tomato', 'SpringGreen', 'Green', 'Yellow', 'Olive',
          'Gold', 'Wheat', 'Orange', 'Gray', 'Red']


def draw(image, bbox, classes, show_conf=False, conf_th=0.0):
    """
    :param image: tensor
    :param bbox: tensor xmin, ymin, xmax, ymax
    """
    keys = list(classes.keys())
    values = list(classes.values())
	
    # 设置字体(包括大小)
    font = ImageFont.truetype('arial.ttf', 10)

    transform = F.ToPILImage()
    image = transform(image)
    draw_image = ImageDraw.Draw(image)

    bbox = np.array(bbox.cpu())

    for b in bbox:
        print(b)
        if show_conf and b[-2] < conf_th:
            continue
        draw_image.rectangle(list(b[:4]), outline=colors[int(b[-1])], width=3)
        if show_conf:
            draw_image.text(list(b[:2] + 5), keys[values.index(int(b[-1]))] + ' {:.2f}'.format(b[-2]),
                            fill=colors[int(b[-1])], font=font)
        else:
            draw_image.text(list(b[:2] + 5), keys[values.index(int(b[-1]))],
                            fill=colors[int(b[-1])], font=font)

    plt.figure()
    plt.imshow(image)
    plt.show()

模型

darknet.py

根据论文的描述,首先构建Darknet的Backbone部分

下图是Backbone的基本结构请添加图片描述,源站可能有防盗链机制,建议将图片保存下来直接上传
请添加图片描述
请添加图片描述
下文指出了darknet全部使用LeakyReLU并且参数为0.1,除了最后的全连接层不需要激活函数。

在这里插入图片描述

# 路径: ./model/darknet.py
import torch.nn as nn


def conv(in_ch, out_ch, k_size=3, stride=1, padding=1):
    return nn.Sequential(
        nn.Conv2d(in_ch, out_ch, k_size, stride, padding, bias=False),
        nn.LeakyReLU(0.1)
    )


def make_layer(param):
    layers = []
    if not isinstance(param[0], list):
        param = [param]
    for p in param:
        layers.append(conv(*p))
    return nn.Sequential(*layers)


class Block(nn.Module):
    def __init__(self, param, use_pool=True):
        super(Block, self).__init__()

        self.conv = make_layer(param)
        self.pool = nn.MaxPool2d(2)
        self.use_pool = use_pool

    def forward(self, x):
        x = self.conv(x)
        if self.use_pool:
            x = self.pool(x)
        return x


class DarkNet(nn.Module):
    def __init__(self):
        super(DarkNet, self).__init__()
		
        self.conv1 = Block([[3, 64, 7, 2, 3]])
        self.conv2 = Block([[64, 192, 3, 1, 1]])
        self.conv3 = Block([[192, 128, 1, 1, 0],
                            [128, 256, 3, 1, 1],
                            [256, 256, 1, 1, 0],
                            [256, 512, 3, 1, 1]])
        self.conv4 = Block([[512, 256, 1, 1, 0],
                            [256, 512, 3, 1, 1],
                            [512, 256, 1, 1, 0],
                            [256, 512, 3, 1, 1],
                            [512, 256, 1, 1, 0],
                            [256, 512, 3, 1, 1],
                            [512, 256, 1, 1, 0],
                            [256, 512, 3, 1, 1],
                            [512, 512, 1, 1, 0],
                            [512, 1024, 3, 1, 1]])
        self.conv5 = Block([[1024, 512, 1, 1, 0],
                            [512, 1024, 3, 1, 1],
                            [1024, 512, 1, 1, 0],
                            [512, 1024, 3, 1, 1],
                            [1024, 1024, 3, 1, 1],
                            [1024, 1024, 3, 2, 1]], False)
        self.conv6 = Block([[1024, 1024, 3, 1, 1],
                            [1024, 1024, 3, 1, 1]], False)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.conv6(x)
        return x


if __name__ == "__main__":
    import torch
    x = torch.randn([1, 3, 448, 448])

    net = DarkNet()
    print(net)
    out = net(x)
    print(out.size())
resnet.py

在这里插入图片描述

作者还说明了Darknet需要预训练,由于在ImageNet上进行预训练耗时过长,我选择使用修改过的Resnet50作为Backbone并使用Pytorch官方的预训练参数。

# 路径: ./model/resnet.py
import torch
from torchvision.models.resnet import ResNet, Bottleneck


"""
通过继承Pytorch的ResNet代码,重写其中的_forward_impl来去除最后的avgpool与fc层
此外我将Resnet50原有的layer4省略并额外增加了两个maxpool层,使得Resnet输出的特征图与Darknet一致
均为[1024, 7, 7]
"""
class ResNet_(ResNet):
    def __init__(self, block, layers):
        super(ResNet_, self).__init__(block=block, layers=layers)

    def _forward_impl(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.maxpool(x)
        x = self.layer3(x)
        x = self.maxpool(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)


def _resnet(block, layers, pretrained):
    model = ResNet_(block, layers)
    if pretrained is not None:
        state_dict = torch.load(pretrained)
        model.load_state_dict(state_dict)
    return model


def resnet_1024ch(pretrained=None) -> ResNet:
    resnet = _resnet(Bottleneck, [3, 4, 6, 3], pretrained)
    return resnet


if __name__ == '__main__':
    x = torch.randn([1, 3, 448, 448])
    net = resnet_1024ch('resnet50-19c8e357.pth')
    print(net)

    y = net(x)
    print(y.size())
yolo.py

接下来复现yolo v1的目标检测核心部分。
在这里插入图片描述
在这里插入图片描述
​ Yolo v1的目标检测思路是将一张图像分割成S × S个grid cell,每个grid cell都会预测B个bbox,同时每个grid cell都会预测一个类别,作为其预测出的B个bbox的共同类别。如果一个目标的中心坐标位于某个grid cell中,那么这个grid cell就负责预测这个目标。一个bbox由x,y,w,h,conf五个参数表示,分别是中心X坐标,中心Y坐标(这里的x、y是相对于grid cell的x、y),宽,高,object置信度。其中这五个参数都是百分比的表示形式,x、y需要除以grid cell的宽、高;w、h则需要除以整张图片的宽、高;对于object conf,如果没有对应bbox则label为0,如果有对应bbox则label为预测出的bbox与对应bbox的IoU交并比。而且在测试时最终conf=object conf * conditional class probabilities。最终结果是xywh都在0到1之间。负责预测类别的tensor长度应为class个数。模型结构上最后一个全连接层的大小则是S × S × ( B * 5 + C )。
在这里插入图片描述
​ 最后预测bbox的全连接层之间有Dropout层,rate=0.5
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

损失函数上,YOLO v1以均方损失为为主,在部分参数的损失计算上使用了一些技巧。如w、h计算中就先开根号使得小目标的损失更为显著,加入λcoord=5和λnoobj=0.5以增加xy的损失值(0到1中开根号会让值变大),降低框出背景bbox的损失以适应该类bbox较多的状况。复现过程中xywhc的五个损失值都直接计算、class的损失值由CorssEntropyLoss函数直接计算。文中也定义了预测的bbox与label中的bbox匹配方法——IoU匹配。以IoU较大的一对bbox确定对应关系,而且这个过程取走一对之后不能重复取。只有对应的bbox才能计算x y w h conf各类损失。如果一个预测bbox没有对应的label bbox,那么就认为其标记到了背景,只计算conf。class的损失是有负责label bbox的grid cell需要计算的。因此其计算次数只是S2次。

这里需注意计算损失时如果没有使用pytorch中官方定义的损失函数,那么需要先通过sigmoid函数将模型输出限制到0~1之间才能进行直接的损失计算。

yolo.py的代码了yolo模型、yolo损失计算、yolo后处理三个部分。对前两个模块,yolo.py都有相应的测试代码,而后处理部分需要在test.py中测试,其不会被yolo.py内部调用。

# 路径: ./model/yolo.py
import numpy as np

from model.darknet import DarkNet
from model.resnet import resnet_1024ch

import torch
import torch.nn as nn
import torchvision


# yolo模型
class yolo(nn.Module):
    def __init__(self, s, cell_out_ch, backbone_name, pretrain=None):
        """
        return: [s, s, cell_out_ch]
        """

        super(yolo, self).__init__()

        self.s = s
        self.backbone = None
        self.conv = None
        if backbone_name == 'darknet':
            self.backbone = DarkNet()
        elif backbone_name == 'resnet':
            self.backbone = resnet_1024ch(pretrained=pretrain)
        self.backbone_name = backbone_name

        assert self.backbone is not None, 'Wrong backbone name'

        self.fc = nn.Sequential(
            nn.Linear(1024 * s * s, 4096),
            nn.LeakyReLU(0.1),
            nn.Dropout(0.5),
            nn.Linear(4096, s * s * cell_out_ch)
        )

    def forward(self, x):
        batch_size = x.size(0)
        x = self.backbone(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        x = x.view(batch_size, self.s ** 2, -1)
        return x


# yolo损失计算
class yolo_loss:
    def __init__(self, device, s, b, image_size, num_classes):
        self.device = device
        self.s = s
        self.b = b
        self.image_size = image_size
        self.num_classes = num_classes
        self.batch_size = 0

    def __call__(self, input, target):
        """
        :param input: (yolo net output)
                      tensor[s, s, b*5 + n_class] bbox: b * (c_x, c_y, w, h, obj_conf), class1_p, class2_p.. %
        :param target: (dataset) tensor[n_bbox] bbox: x_min, ymin, xmax, ymax, class
        :return: loss tensor

        grid type: [[bbox, ..], [], ..] -> bbox_in_grid: c_x(%), c_y(%), w(%), h(%), class(int)

        target to grid type
        if s = 7 -> grid idx: 1 -> 49
        由于没有使用PyTorch的损失函数,所以需要先分离不同的batch分别计算损失
        """
        self.batch_size = input.size(0)
		
        # label预处理
        target = [self.label_direct2grid(target[i]) for i in range(self.batch_size)]

        # IoU 匹配predictor和label
        # 以Predictor为基准,每个Predictor都有且仅有一个需要负责的Target(前提是Predictor所在Grid Cell有Target中心位于此)
        # x, y, w, h, c
        match = []
        conf = []
        for i in range(self.batch_size):
            m, c = self.match_pred_target(input[i], target[i])
            match.append(m)
            conf.append(c)

        loss = torch.zeros([self.batch_size], dtype=torch.float, device=self.device)
        xy_loss = torch.zeros_like(loss)
        wh_loss = torch.zeros_like(loss)
        conf_loss = torch.zeros_like(loss)
        class_loss = torch.zeros_like(loss)
        for i in range(self.batch_size):
            loss[i], xy_loss[i], wh_loss[i], conf_loss[i], class_loss[i] = \
                self.compute_loss(input[i], target[i], match[i], conf[i])
        return torch.mean(loss), torch.mean(xy_loss), torch.mean(wh_loss), torch.mean(conf_loss), torch.mean(class_loss)

    def label_direct2grid(self, label):
        """
        :param label: dataset type: xmin, ymin, xmax, ymax, class
        :return: label: grid type, if the grid doesn't have object -> put None
        将label转换为c_x, c_y, w, h, conf再根据不同的grid cell分类,并转换成百分比形式
        若一个grid cell中没有label则都用None代替
        """
        output = [None for _ in range(self.s ** 2)]
        size = self.image_size // self.s  # h, w

        n_bbox = label.size(0)
        label_c = torch.zeros_like(label)

        label_c[:, 0] = (label[:, 0] + label[:, 2]) / 2
        label_c[:, 1] = (label[:, 1] + label[:, 3]) / 2
        label_c[:, 2] = abs(label[:, 0] - label[:, 2])
        label_c[:, 3] = abs(label[:, 1] - label[:, 3])
        label_c[:, 4] = label[:, 4]

        idx_x = [int(label_c[i][0]) // size for i in range(n_bbox)]
        idx_y = [int(label_c[i][1]) // size for i in range(n_bbox)]

        label_c[:, 0] = torch.div(torch.fmod(label_c[:, 0], size), size)
        label_c[:, 1] = torch.div(torch.fmod(label_c[:, 1], size), size)
        label_c[:, 2] = torch.div(label_c[:, 2], self.image_size)
        label_c[:, 3] = torch.div(label_c[:, 3], self.image_size)

        for i in range(n_bbox):
            idx = idx_y[i] * self.s + idx_x[i]
            if output[idx] is None:
                output[idx] = torch.unsqueeze(label_c[i], dim=0)
            else:
                output[idx] = torch.cat([output[idx], torch.unsqueeze(label_c[i], dim=0)], dim=0)
        return output

    def match_pred_target(self, input, target):
        match = []
        conf = []
        with torch.no_grad():
            input_bbox = input[:, :self.b * 5].reshape(-1, self.b, 5)
            ious = [match_get_iou(input_bbox[i], target[i], self.s, i)
                    for i in range(self.s ** 2)]
            for iou in ious:
                if iou is None:
                    match.append(None)
                    conf.append(None)
                else:
                    keep = np.ones([len(iou[0])], dtype=bool)
                    m = []
                    c = []
                    for i in range(self.b):
                        if np.any(keep) == False:
                            break
                        idx = np.argmax(iou[i][keep])
                        np_max = np.max(iou[i][keep])
                        m.append(np.argwhere(iou[i] == np_max).tolist()[0][0])
                        c.append(np.max(iou[i][keep]))
                        keep[idx] = 0
                    match.append(m)
                    conf.append(c)
        return match, conf

    def compute_loss(self, input, target, match, conf):
        # 计算损失
        ce_loss = nn.CrossEntropyLoss()

        input_bbox = input[:, :self.b * 5].reshape(-1, self.b, 5)
        input_class = input[:, self.b * 5:].reshape(-1, self.num_classes)

        input_bbox = torch.sigmoid(input_bbox)
        loss = torch.zeros([self.s ** 2], dtype=torch.float, device=self.device)
        xy_loss = torch.zeros_like(loss)
        wh_loss = torch.zeros_like(loss)
        conf_loss = torch.zeros_like(loss)
        class_loss = torch.zeros_like(loss)
        # 不同grid cell分别计算再求和
        for i in range(self.s ** 2):
            # 0 xy_loss, 1 wh_loss, 2 conf_loss, 3 class_loss
            l = torch.zeros([4], dtype=torch.float, device=self.device)
            # Neg
            if target[i] is None:
                # λ_noobj = 0.5
                obj_conf_target = torch.zeros([self.b], dtype=torch.float, device=self.device)
                l[2] = torch.sum(torch.mul(0.5, torch.pow(input_bbox[i, :, 4] - obj_conf_target, 2)))
            else:
                # λ_coord = 5
                l[0] = torch.mul(5, torch.sum(torch.pow(input_bbox[i, :, 0] - target[i][match[i], 0], 2) +
                                              torch.pow(input_bbox[i, :, 1] - target[i][match[i], 1], 2)))

                l[1] = torch.mul(5, torch.sum(torch.pow(torch.sqrt(input_bbox[i, :, 2]) -
                                                        torch.sqrt(target[i][match[i], 2]), 2) +
                                              torch.pow(torch.sqrt(input_bbox[i, :, 3]) -
                                                        torch.sqrt(target[i][match[i], 3]), 2)))
                obj_conf_target = torch.tensor(conf[i], dtype=torch.float, device=self.device)
                l[2] = torch.sum(torch.pow(input_bbox[i, :, 4] - obj_conf_target, 2))

                l[3] = ce_loss(input_class[i].unsqueeze(dim=0).repeat(target[i].size(0), 1),
                               target[i][:, 4].long())
            loss[i] = torch.sum(l)
            xy_loss[i] = torch.sum(l[0])
            wh_loss[i] = torch.sum(l[1])
            conf_loss[i] = torch.sum(l[2])
            class_loss[i] = torch.sum(l[3])
        return torch.sum(loss), torch.sum(xy_loss), torch.sum(wh_loss), torch.sum(conf_loss), torch.sum(class_loss)


def cxcywh2xyxy(bbox):
    """
    :param bbox: [bbox, bbox, ..] tensor c_x(%), c_y(%), w(%), h(%), c
    """
    bbox[:, 0] = bbox[:, 0] - bbox[:, 2] / 2
    bbox[:, 1] = bbox[:, 1] - bbox[:, 3] / 2
    bbox[:, 2] = bbox[:, 0] + bbox[:, 2]
    bbox[:, 3] = bbox[:, 1] + bbox[:, 3]
    return bbox


def match_get_iou(bbox1, bbox2, s, idx):
    """
    :param bbox1: [bbox, bbox, ..] tensor c_x(%), c_y(%), w(%), h(%), c
    :return:
    """

    if bbox1 is None or bbox2 is None:
        return None

    bbox1 = np.array(bbox1.cpu())
    bbox2 = np.array(bbox2.cpu())
	
    # c_x, c_y转换为对整张图片的百分比
    bbox1[:, 0] = bbox1[:, 0] / s
    bbox1[:, 1] = bbox1[:, 1] / s
    bbox2[:, 0] = bbox2[:, 0] / s
    bbox2[:, 1] = bbox2[:, 1] / s

    # c_x, c_y加上grid cell左上角左边变成完整坐标
    grid_pos = [(j / s, i / s) for i in range(s) for j in range(s)]
    bbox1[:, 0] = bbox1[:, 0] + grid_pos[idx][0]
    bbox1[:, 1] = bbox1[:, 1] + grid_pos[idx][1]
    bbox2[:, 0] = bbox2[:, 0] + grid_pos[idx][0]
    bbox2[:, 1] = bbox2[:, 1] + grid_pos[idx][1]

    bbox1 = cxcywh2xyxy(bbox1)
    bbox2 = cxcywh2xyxy(bbox2)

    # %
    return get_iou(bbox1, bbox2)


def get_iou(bbox1, bbox2):
    """
    :param bbox1: [bbox, bbox, ..] tensor xmin ymin xmax ymax
    :param bbox2:
    :return: area:
    """

    s1 = abs(bbox1[:, 2] - bbox1[:, 0]) * abs(bbox1[:, 3] - bbox1[:, 1])
    s2 = abs(bbox2[:, 2] - bbox2[:, 0]) * abs(bbox2[:, 3] - bbox2[:, 1])

    ious = []
    for i in range(bbox1.shape[0]):
        xmin = np.maximum(bbox1[i, 0], bbox2[:, 0])
        ymin = np.maximum(bbox1[i, 1], bbox2[:, 1])
        xmax = np.minimum(bbox1[i, 2], bbox2[:, 2])
        ymax = np.minimum(bbox1[i, 3], bbox2[:, 3])

        in_w = np.maximum(xmax - xmin, 0)
        in_h = np.maximum(ymax - ymin, 0)

        in_s = in_w * in_h

        iou = in_s / (s1[i] + s2 - in_s)
        ious.append(iou)
    ious = np.array(ious)
    return ious


def nms(bbox, conf_th, iou_th):
    bbox = np.array(bbox.cpu())

    bbox[:, 4] = bbox[:, 4] * bbox[:, 5]

    bbox = bbox[bbox[:, 4] > conf_th]
    order = np.argsort(-bbox[:, 4])

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        iou = get_iou(np.array([bbox[i]]), bbox[order[1:]])[0]
        inds = np.where(iou <= iou_th)[0]
        order = order[inds + 1]
    return bbox[keep]


# yolo后处理
def output_process(output, image_size, s, b, conf_th, iou_th):
    """
    输入是包含batch的模型输出
    :return output: list[], bbox: xmin, ymin, xmax, ymax, obj_conf, classes_conf, classes
    """
    batch_size = output.size(0)
    size = image_size // s

    output = torch.sigmoid(output)

    # Get Class
    # 将class conf依次添加到bbox中
    classes_conf, classes = torch.max(output[:, :, b * 5:], dim=2)
    classes = classes.unsqueeze(dim=2).repeat(1, 1, 2).unsqueeze(dim=3)
    classes_conf = classes_conf.unsqueeze(dim=2).repeat(1, 1, 2).unsqueeze(dim=3)
    bbox = output[:, :, :b * 5].reshape(batch_size, -1, b, 5)

    bbox = torch.cat([bbox, classes_conf, classes], dim=3)

    # To Direct
    # 百分比形式转直接表示
    bbox[:, :, :, [0, 1]] = bbox[:, :, :, [0, 1]] * size
    bbox[:, :, :, [2, 3]] = bbox[:, :, :, [2, 3]] * image_size
	
    # 添加grid cell坐标
    grid_pos = [(j * image_size // s, i * image_size // s) for i in range(s) for j in range(s)]

    def to_direct(bbox):
        for i in range(s ** 2):
            bbox[i, :, 0] = bbox[i, :, 0] + grid_pos[i][0]
            bbox[i, :, 1] = bbox[i, :, 1] + grid_pos[i][1]
        return bbox

    bbox_direct = torch.stack([to_direct(b) for b in bbox])
    bbox_direct = bbox_direct.reshape(batch_size, -1, 7)

    # cxcywh to xyxy
    bbox_direct[:, :, 0] = bbox_direct[:, :, 0] - bbox_direct[:, :, 2] / 2
    bbox_direct[:, :, 1] = bbox_direct[:, :, 1] - bbox_direct[:, :, 3] / 2
    bbox_direct[:, :, 2] = bbox_direct[:, :, 0] + bbox_direct[:, :, 2]
    bbox_direct[:, :, 3] = bbox_direct[:, :, 1] + bbox_direct[:, :, 3]

    bbox_direct[:, :, 0] = torch.maximum(bbox_direct[:, :, 0], torch.zeros(1))
    bbox_direct[:, :, 1] = torch.maximum(bbox_direct[:, :, 1], torch.zeros(1))
    bbox_direct[:, :, 2] = torch.minimum(bbox_direct[:, :, 2], torch.tensor([image_size]))
    bbox_direct[:, :, 3] = torch.minimum(bbox_direct[:, :, 3], torch.tensor([image_size]))
	
    # 整合不同batch中的bbox
    bbox = [torch.tensor(nms(b, conf_th, iou_th)) for b in bbox_direct]
    bbox = torch.stack(bbox)
    return bbox


if __name__ == "__main__":
    import torch

    # Test yolo
    x = torch.randn([1, 3, 448, 448])

    # B * 5 + n_classes
    net = yolo(7, 2 * 5 + 20, 'resnet', pretrain=None)
    # net = yolo(7, 2 * 5 + 20, 'darknet', pretrain=None)
    print(net)
    out = net(x)
    print(out)
    print(out.size())

    # Test yolo_loss
    # 测试时假设 s=2, class=2
    s = 2
    b = 2
    image_size = 448  # h, w
    input = torch.tensor([[[0.45, 0.24, 0.22, 0.3, 0.35, 0.54, 0.66, 0.7, 0.8, 0.8, 0.17, 0.9],
                           [0.37, 0.25, 0.5, 0.3, 0.36, 0.14, 0.27, 0.26, 0.33, 0.36, 0.13, 0.9],
                           [0.12, 0.8, 0.26, 0.74, 0.8, 0.13, 0.83, 0.6, 0.75, 0.87, 0.75, 0.24],
                           [0.1, 0.27, 0.24, 0.37, 0.34, 0.15, 0.26, 0.27, 0.37, 0.34, 0.16, 0.93]]])
    target = [torch.tensor([[200, 200, 353, 300, 1],
                            [220, 230, 353, 300, 1],
                            [15, 330, 200, 400, 0],
                            [100, 50, 198, 223, 1],
                            [30, 60, 150, 240, 1]], dtype=torch.float)]

    criterion = yolo_loss('cpu', 2, 2, image_size, 2)
    loss = criterion(input, target)
    print(loss)

scheduler.py

在这里插入图片描述

论文作者也给出了他的训练参数。我们首先复现学习率的调整方式。原文分成三段学习率,每一段都有不同的保持epochs长度。此外,我额外增加了热身训练阶段。

# 路径: ./scheduler.py
from torch.optim.lr_scheduler import _LRScheduler


class Scheduler(_LRScheduler):
    def __init__(self, optimizer, step_warm_ep, lr_start, step_1_lr, step_1_ep,
                 step_2_lr, step_2_ep, step_3_lr, step_3_ep, last_epoch=-1):
        self.optimizer = optimizer
        self.lr_start = lr_start
        self.step_warm_ep = step_warm_ep
        self.step_1_lr = step_1_lr
        self.step_1_ep = step_1_ep
        self.step_2_lr = step_2_lr
        self.step_2_ep = step_2_ep
        self.step_3_lr = step_3_lr
        self.step_3_ep = step_3_ep
        self.last_epoch = last_epoch

        super(Scheduler, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch == 0:
            return [self.lr_start for _ in self.optimizer.param_groups]
        lr = self._compute_lr_from_epoch()

        return [lr for _ in self.optimizer.param_groups]

    def _get_closed_form_lr(self):
        return self.base_lrs

    def _compute_lr_from_epoch(self):
        if self.last_epoch < self.step_warm_ep:
            lr = ((self.step_1_lr - self.lr_start)/self.step_warm_ep) * self.last_epoch + self.lr_start
        elif self.last_epoch < self.step_warm_ep + self.step_1_ep:
            lr = self.step_1_lr
        elif self.last_epoch < self.step_warm_ep + self.step_1_ep + self.step_2_ep:
            lr = self.step_2_lr
        elif self.last_epoch < self.step_warm_ep + self.step_1_ep + self.step_2_ep + self.step_3_ep:
            lr = self.step_3_lr
        else:
            lr = self.step_3_lr
        return lr


if __name__ == '__main__':
    import torch.nn as nn
    import torch.optim as optim

    import numpy as np
    import matplotlib.pyplot as plt

    import warnings
    warnings.filterwarnings('ignore')

    batch_size = 16
    epoch = 135
    scheduler_params = {
        'lr_start': 1e-3,
        'step_warm_ep': 10,
        'step_1_lr': 1e-2,
        'step_1_ep': 75,
        'step_2_lr': 1e-3,
        'step_2_ep': 30,
        'step_3_lr': 1e-4,
        'step_3_ep': 20
    }

    model = nn.Sequential(
        nn.Linear(1, 10),
        nn.Linear(10, 1)
    )
    optimizer = optim.SGD(model.parameters(), lr=scheduler_params['lr_start'])
    scheduler = Scheduler(optimizer, **scheduler_params)

    lrs = []
    for _ in range(epoch):
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()
    print(lrs)

    lrs = np.array(lrs)

    # 使用plt可视化学习率
    plt.figure()
    plt.plot(lrs)
    plt.show()

train.py

训练参数主要参考了论文原文。由于batch size较小,学习率都等比减小。优化器选择了SGD并带有momentum与weight_decay。代码中有冻结Backbone的选项。

root0712 存储了VOC07/12两年数据集的根目录

model_root 存储模型路径

backbone 'resnet’则代表使用修改后的resnet;'darknet’则代表使用原文中的darknet

pretrain None则不使用预训练模型;或直接输入预训练权重地址

with_amp 混合精度选项

transforms 需要使用代码中定义的transforms而不是PyTorch直接给出的transforms

start_epoch 中断训练重启的开始epoch

epoch 总epoch数

freeze_backbone_till 若为-1则不冻结backbone,其他数则会冻结backbone直到该epoch后解冻

# 路径: ./train.py
from dataset.data import VOC0712Dataset
from dataset.transform import *
from model.yolo import yolo, yolo_loss
from scheduler import Scheduler

import torch
from torch.utils.data import DataLoader
from torch import optim
from torch.cuda.amp import autocast, GradScaler

from tqdm import tqdm
import pandas as pd
import json
import os

import warnings


class CFG:
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    root0712 = [r'F:\AI\Dataset\VOC2007\VOCdevkit\VOC2007', r'F:\AI\Dataset\VOC2012\VOCdevkit\VOC2012']
    class_path = r'./dataset/classes.json'
    model_root = r'./log/ex7'
    # 若model_path是一个指向权重文件的str路径,那么会将模型传入指定模型权重
    model_path = None
	# 这里没有编写创建文件夹的代码,直接运行需要手动将文件夹创建好
    
    backbone = 'resnet'
    pretrain = 'model/resnet50-19c8e357.pth'
    # 混合精度可选,若为False则采用常规精度
    with_amp = True
    S = 7
    B = 2
    image_size = 448

    transforms = Compose([
        ToTensor(),
        RandomHorizontalFlip(0.5),
        Resize(448, keep_ratio=False)
    ])

    start_epoch = 0
    epoch = 135
    batch_size = 16
    num_workers = 2
	
    # freeze_backbone_till = -1 则不冻结
    freeze_backbone_till = 30

    scheduler_params = {
        'lr_start': 1e-3 / 4,
        'step_warm_ep': 10,
        'step_1_lr': 1e-2 / 4,
        'step_1_ep': 75,
        'step_2_lr': 1e-3 / 4,
        'step_2_ep': 40,
        'step_3_lr': 1e-4 / 4,
        'step_3_ep': 10
    }

    momentum = 0.9
    weight_decay = 0.0005


def collate_fn(batch):
    return tuple(zip(*batch))


class AverageMeter:
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def train():
    device = torch.device(CFG.device)
    print('Train:\nDevice:{}'.format(device))

    with open(CFG.class_path, 'r') as f:
        json_str = f.read()
        classes = json.loads(json_str)
        CFG.num_classes = len(classes)

    train_ds = VOC0712Dataset(CFG.root0712, CFG.class_path, CFG.transforms, 'train')
    test_ds = VOC0712Dataset(CFG.root0712, CFG.class_path, CFG.transforms, 'test')

    train_dl = DataLoader(train_ds, batch_size=CFG.batch_size, shuffle=True,
                          num_workers=CFG.num_workers, collate_fn=collate_fn)
    test_dl = DataLoader(test_ds, batch_size=CFG.batch_size, shuffle=False,
                         num_workers=CFG.num_workers, collate_fn=collate_fn)

    yolo_net = yolo(s=CFG.S, cell_out_ch=CFG.B * 5 + CFG.num_classes, backbone_name=CFG.backbone, pretrain=CFG.pretrain)
    yolo_net.to(device)

    if CFG.model_path is not None:
        yolo_net.load_state_dict(torch.load(CFG.model_path))

    if CFG.freeze_backbone_till != -1:
        print('Freeze Backbone')
        for param in yolo_net.backbone.parameters():
            param.requires_grad_(False)

    param = [p for p in yolo_net.parameters() if p.requires_grad]
    optimizer = optim.SGD(param, lr=CFG.scheduler_params['lr_start'],
                          momentum=CFG.momentum, weight_decay=CFG.weight_decay)
    criterion = yolo_loss(CFG.device, CFG.S, CFG.B, CFG.image_size, len(train_ds.classes))
    scheduler = Scheduler(optimizer, **CFG.scheduler_params)
    scaler = GradScaler()

    for _ in range(CFG.start_epoch):
        scheduler.step()

    best_train_loss = 1e+9
    train_losses = []
    test_losses = []
    lrs = []
    for epoch in range(CFG.start_epoch, CFG.epoch):
        if CFG.freeze_backbone_till != -1 and epoch >= CFG.freeze_backbone_till:
            print('Unfreeze Backbone')
            for param in yolo_net.backbone.parameters():
                param.requires_grad_(True)
            CFG.freeze_backbone_till = -1
        # Train
        yolo_net.train()
        loss_score = AverageMeter()
        dl = tqdm(train_dl, total=len(train_dl))
        for images, labels in dl:
            batch_size = len(labels)

            images = torch.stack(images)
            images = images.to(device)
            labels = [label.to(device) for label in labels]

            optimizer.zero_grad()

            if CFG.with_amp:
                with autocast():
                    outputs = yolo_net(images)
                    loss, xy_loss, wh_loss, conf_loss, class_loss = criterion(outputs, labels)

                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                outputs = yolo_net(images)
                loss, xy_loss, wh_loss, conf_loss, class_loss = criterion(outputs, labels)

                loss.backward()
                optimizer.step()

            loss_score.update(loss.detach().item(), batch_size)
            dl.set_postfix(Mode='Train', AvgLoss=loss_score.avg, Loss=loss.detach().item(),
                           Epoch=epoch, LR=optimizer.param_groups[0]['lr'])
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()

        train_losses.append(loss_score.avg)
        print('Train Loss: {:.4f}'.format(loss_score.avg))

        if best_train_loss > loss_score.avg:
            print('Save yolo_net to {}'.format(os.path.join(CFG.model_root, 'yolo.pth')))
            torch.save(yolo_net.state_dict(), os.path.join(CFG.model_root, 'yolo.pth'))
            best_train_loss = loss_score.avg

        loss_score.reset()
        with torch.no_grad():
            # Test
            yolo_net.eval()
            dl = tqdm(test_dl, total=len(test_dl))
            for images, labels in dl:
                batch_size = len(labels)

                images = torch.stack(images)
                images = images.to(device)
                labels = [label.to(device) for label in labels]

                outputs = yolo_net(images)
                loss, xy_loss, wh_loss, conf_loss, class_loss = criterion(outputs, labels)

                loss_score.update(loss.detach().item(), batch_size)
                dl.set_postfix(Mode='Test', AvgLoss=loss_score.avg, Loss=loss.detach().item(), Epoch=epoch)
        test_losses.append(loss_score.avg)
        print('Test Loss: {:.4f}'.format(loss_score.avg))

        df = pd.DataFrame({'Train Loss': train_losses, 'Test Loss': test_losses, 'LR': lrs})
        df.to_csv(os.path.join(CFG.model_root, 'result.csv'), index=True)


if __name__ == '__main__':
    warnings.filterwarnings('ignore')
    train()

voc_eval.py

测试时计算mAP的代码

# 路径: ./voc_eval.py
"""Adapted from:
    @longcw faster_rcnn_pytorch: https://github.com/longcw/faster_rcnn_pytorch
    @rbgirshick py-faster-rcnn https://github.com/rbgirshick/py-faster-rcnn
    Licensed under The MIT License [see LICENSE for details]
    这个文件的代码有改动,我删除了缓存的代码,因此也不需要传入缓存文件夹
"""

import xml.etree.ElementTree as ET
import numpy as np
import os


def voc_ap(rec, prec, use_07_metric=False):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([0.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap


def parse_rec(filename):
    """ Parse a PASCAL VOC xml file """
    tree = ET.parse(filename)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        obj_struct['name'] = obj.find('name').text
        obj_struct['pose'] = obj.find('pose').text
        obj_struct['truncated'] = int(obj.find('truncated').text)
        obj_struct['difficult'] = int(obj.find('difficult').text)
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(bbox.find('xmin').text),
                              int(bbox.find('ymin').text),
                              int(bbox.find('xmax').text),
                              int(bbox.find('ymax').text)]
        objects.append(obj_struct)

    return objects


def voc_eval(detpath,
             annopath,
             imagesetfile,
             classname,
             ovthresh=0.5,
             use_07_metric=False):
    """rec, prec, ap = voc_eval(detpath,
                                annopath,
                                imagesetfile,
                                classname,
                                [ovthresh],
                                [use_07_metric])
    Top level function that does the PASCAL VOC evaluation.
    detpath: Path to detections
        detpath.format(classname) should produce the detection results file.
    annopath: Path to annotations
        annopath.format(imagename) should be the xml annotations file.
    imagesetfile: Text file containing the list of images, one image per line.
    classname: Category name (duh)
    cachedir: Directory for caching the annotations
    [ovthresh]: Overlap threshold (default = 0.5)
    [use_07_metric]: Whether to use VOC07's 11 point AP computation
        (default False)
    """
    # assumes detections are in detpath.format(classname)
    # assumes annotations are in annopath.format(imagename)
    # assumes imagesetfile is a text file with each line an image name
    # cachedir caches the annotations in a pickle file

    # read list of images
    with open(imagesetfile, 'r') as f:
        lines = f.readlines()
    imagenames = [x.strip() for x in lines]

    # load annots
    recs = {}
    for i, imagename in enumerate(imagenames):
        recs[imagename] = parse_rec(annopath.format(imagename))
        # if i % 100 == 0:
        #     print('Reading annotation for {:d}/{:d}'.format(
        #         i + 1, len(imagenames)))

    # extract gt objects for this class
    class_recs = {}
    npos = 0
    for imagename in imagenames:
        R = [obj for obj in recs[imagename] if obj['name'] == classname]
        bbox = np.array([x['bbox'] for x in R])
        difficult = np.array([x['difficult'] for x in R]).astype(np.bool)
        det = [False] * len(R)
        npos = npos + sum(~difficult)
        class_recs[imagename] = {'bbox': bbox,
                                 'difficult': difficult,
                                 'det': det}

    # read dets
    detfile = detpath.format(classname)
    with open(detfile, 'r') as f:
        lines = f.readlines()

    splitlines = [x.strip().split(' ') for x in lines]
    image_ids = [x[0] for x in splitlines]
    confidence = np.array([float(x[1]) for x in splitlines])
    BB = np.array([[float(z) for z in x[2:]] for x in splitlines])

    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    sorted_scores = np.sort(-confidence)
    BB = BB[sorted_ind, :]
    image_ids = [image_ids[x] for x in sorted_ind]

    # go down dets and mark TPs and FPs
    nd = len(image_ids)
    tp = np.zeros(nd)
    fp = np.zeros(nd)
    for d in range(nd):
        R = class_recs[image_ids[d]]
        bb = BB[d, :].astype(float)
        ovmax = -np.inf
        BBGT = R['bbox'].astype(float)

        if BBGT.size > 0:
            # compute overlaps
            # intersection
            ixmin = np.maximum(BBGT[:, 0], bb[0])
            iymin = np.maximum(BBGT[:, 1], bb[1])
            ixmax = np.minimum(BBGT[:, 2], bb[2])
            iymax = np.minimum(BBGT[:, 3], bb[3])
            iw = np.maximum(ixmax - ixmin + 1., 0.)
            ih = np.maximum(iymax - iymin + 1., 0.)
            inters = iw * ih

            # union
            uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) +
                   (BBGT[:, 2] - BBGT[:, 0] + 1.) *
                   (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters)

            overlaps = inters / uni
            ovmax = np.max(overlaps)
            jmax = np.argmax(overlaps)
        if ovmax > ovthresh:
            if not R['difficult'][jmax]:
                if not R['det'][jmax]:
                    tp[d] = 1.
                    R['det'][jmax] = 1
                else:
                    fp[d] = 1.
        else:
            fp[d] = 1.

    # compute precision recall
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)
    rec = tp / float(npos)
    # avoid divide by zero in case the first detection matches a difficult
    # ground truth
    prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    ap = voc_ap(rec, prec, use_07_metric)

    return rec, prec, ap

test.py

# 路径: ./test.py
from dataset.data import VOC0712Dataset, Compose, ToTensor, Resize
from dataset.draw_bbox import draw
from model.yolo import yolo, output_process
from voc_eval import voc_eval

import torch

from tqdm import tqdm
import json


class CFG:
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    root = r'F:\AI\Dataset\VOC2007\VOCdevkit\VOC2007'
    class_path = r'dataset/classes.json'
    model_path = r'log/ex7/yolo.pth'
    # 这里也要手动新建det文件夹,用于保存每个class的目标情况
    detpath = r'det\{}.txt'
    annopath = r'F:\AI\Dataset\VOC2007\VOCdevkit\VOC2007\Annotations\{}.xml'
    imagesetfile = r'F:\AI\Dataset\VOC2007\VOCdevkit\VOC2007\ImageSets\Main\test.txt'
    classname = None

    test_range = None
    show_image = False
    get_ap = True

    backbone = 'resnet'
    S = 7
    B = 2
    image_size = 448

    get_info = True
    transforms = Compose([
        ToTensor(),
        Resize(448, keep_ratio=False)
    ])

    num_classes = 0

    conf_th = 0.2
    iou_th = 0.5


def test():
    device = torch.device(CFG.device)
    print('Test:\nDevice:{}'.format(device))

    dataset = VOC0712Dataset(CFG.root, CFG.class_path, CFG.transforms, 'test',
                             data_range=CFG.test_range, get_info=CFG.get_info)

    with open(CFG.class_path, 'r') as f:
        json_str = f.read()
        classes = json.loads(json_str)
        CFG.classname = list(classes.keys())
        CFG.num_classes = len(CFG.classname)

    yolo_net = yolo(s=CFG.S, cell_out_ch=CFG.B * 5 + CFG.num_classes, backbone_name=CFG.backbone)
    yolo_net.to(device)

    yolo_net.load_state_dict(torch.load(CFG.model_path))

    bboxes = []
    with torch.no_grad():
        for image, label, image_name, input_size in tqdm(dataset):
            image = image.unsqueeze(dim=0)
            image = image.to(device)

            output = yolo_net(image)
            output = output_process(output.cpu(), CFG.image_size, CFG.S, CFG.B, CFG.conf_th, CFG.iou_th)

            if CFG.show_image:
                draw(image.squeeze(dim=0), output.squeeze(dim=0), classes, show_conf=True)
                draw(image.squeeze(dim=0), label, classes, show_conf=True)

            # 还原
            output[:, :, [0, 2]] = output[:, :, [0, 2]] * input_size[0] / CFG.image_size
            output[:, :, [1, 3]] = output[:, :, [1, 3]] * input_size[1] / CFG.image_size

            output = output.squeeze(dim=0).numpy().tolist()
            if len(output) > 0:
                pred = [[image_name, output[i][-3] * output[i][-2]] + output[i][:4] + [int(output[i][-1])]
                        for i in range(len(output))]
                bboxes += pred
    det_list = [[] for _ in range(CFG.num_classes)]
    for b in bboxes:
        det_list[b[-1]].append(b[:-1])

    if CFG.get_ap:
        map = 0
        for idx in range(CFG.num_classes):
            file_path = CFG.detpath.format(CFG.classname[idx])
            txt = '\n'.join([' '.join([str(i) for i in item]) for item in det_list[idx]])
            with open(file_path, 'w') as f:
                f.write(txt)

            rec, prec, ap = voc_eval(CFG.detpath, CFG.annopath, CFG.imagesetfile, CFG.classname[idx])
            print(rec)
            print(prec)
            map += ap
            print(ap)

        map /= CFG.num_classes
        print('mAP', map)


if __name__ == '__main__':
    test()

你可能感兴趣的:(目标检测,深度学习,yolo,YOLO,pytorch,人工智能,深度学习,视觉检测,目标检测,python)