Yolo V1论文笔记
原文地址:
https://arxiv.org/pdf/1506.02640.pdf
源码地址:
https://github.com/motokimura/yolo_v1_pytorch
一句话总结:
YOLO(You Only Look Once)是一种统一的、实时的对象检测方法,通过将对象检测视为一个回归问题来直接从图像中预测边界框和类别概率,实现了快速且准确的对象检测。
网络层最巧妙的设计:
分类网络最后的全连接层,一般连接于一个一维向量,向量的不同位代表不同类别,而这里的输出向量是一个三维的张量(7乘7乘30)。包含了边界框的,中心点坐标,宽高,置信度,类别等信息
在Yolo中,如果一个物体的中心点,落在了某个格子中,那么这个格子将负责预测这个物体。而那些没有物体中心点落进来的格子,则不负责预测任何物体。这个设定就好比该网络在一开始,就将整个图片上的预测任务进行了分工,一共设定7乘7个按照方阵列队的检测人员,每个人员负责检测一个物体,大家的分工界线,就是看被检测物体的中心点落在谁的格子里。当然,是7乘7还是9乘9,是上图中的参数S,可以自己修改,精度和性能会随之有些变化。
(2)30的含义
刚才设定了49个检测人员,那么每个人员负责检测的内容,就是这里的30(注意,30是张量最后一维的长度)。在Yolo v1论文中,30是由(4+1)×2+20得到的。其中4+1是矩形框的中心点坐标(x,y)、长宽(w,h)以及是否属于被检测物体的置信度c;2是一个格子共回归两个矩形框,每个矩形框分别产生5个预测值(每个格子预测矩形框个数,是可调超参数;论文中选择了2个框,当然也可以只预测1个框,具体预测几个矩形框,无非是在计算量和精度之间取一个权衡。如果只预测一个矩形框,计算量会小很多,但是如果训练数据都是小物体,那么网络学习到的框,也会普遍比较小,测试时如果物体较大,那么预测效果就会不理想;如果每个格子多预测几个矩形框,如上文中讲到的,每个矩形框的学习目标会有所分工,有些学习小物体特征,有些学习大物体特征等;在Yolov2、v3中,这个数目都有一定的调整。);20代表预测20个类别。这里有几点需要注意:1. 每个方格(grid) 产生2个预测框,2也是参数,可以调,但是一旦设定为2以后,那么每个方格只产生两个矩形框,最后选定置信度更大的矩形框作为输出,也就是最终每个方格只输出一个预测矩形框。2. 每个方格只能预测一个物体。虽然可以通过调整参数,产生不同的矩形框,但这只能提高矩形框的精度。所以当有很多个物体的中心点落在了同一个格子里,该格子只能预测一个物体。也就是格子数为7乘7时,该网络最多预测49个物体。如上述原文中提及,在强行施加了格点限制以后,每个格点只能输出一个预测结果,所以该算法最大的不足,就是对一些邻近小物体的识别效果不是太好,例如成群结队的小鸟。
文章的观点:
•YOLO框架通过一个单一的神经网络同时预测多个边界框和类别概率,与传统的基于滑动窗口或区域提议的方法相比,YOLO在处理速度和实时性方面具有显著优势。
•该方法在保持高平均精度的同时实现了端到端的训练和实时速度,对于大型物体或边界处的物体可以通过多个单元格进行良好定位,但对于小物体或群体中的小物体检测存在挑战。
•YOLO在跨域泛化能力方面优于其他检测方法,例如从自然图像到艺术作品的泛化,但在精确定位某些物体,特别是小型物体方面仍有提升空间。
相关问题:
🤔Q: YOLO如何实现实时的对象检测?
A: YOLO通过将对象检测问题转化为一个回归问题,使用单个神经网络直接从完整图像中预测边界框和类别概率,无需复杂的处理流程,从而实现了高速处理。
🤔Q: YOLO与传统对象检测方法相比有哪些优势?
A: 相比于基于滑动窗口和区域提议的传统方法,YOLO能够看到整个图像的全局信息,减少了背景误报的数量,同时具有更好的实时性和高泛化能力。
1.更快,
2.全局视野,
3.泛化能力更强
🤔Q: YOLO在检测精度方面存在哪些局限性?
A: YOLO在处理小型物体或群体中的小物体时可能会遇到困难,这是由于其强大的空间约束以及使用相对粗糙的特征进行边界框预测导致的。此外,它对不同尺寸或配置的新对象的泛化能力有限。
关键信息点:
接下来,我将使用提供的信息生成一个思维导图以更直观地展示YOLO对象检测系统的关键概念和特点。
这张思维导图直观地展示了YOLO对象检测系统的核心概念和特点,包括其实时性和准确性、主要优势、存在的局限性,以及技术细节。
Unified Detection(统一检测)
这段内容描述了YOLO(You Only Look Once)对象检测系统的“统一检测”机制,该机制通过单个神经网络集成了对象检测的多个组成部分。以下是对这一段内容的详细解释:
整体分析:
这张图片展示的是YOLO(You Only Look Once)对象检测模型的视觉解释,包含以下几个关键点:
1.S × S网格:图片说明了模型如何将输入的图像划分成\( S \times S \)的网格结构。在这个例子中,图像被划分成7x7的网格(\( S = 7 \))。
2.边界框和置信度:每个网格单元(每个彩色的方块代表一个网格单元)负责预测\( B \)个边界框以及这些框的置信度分数。这个例子中,每个网格单元预测2个边界框(\( B = 2 \))。
3.类别概率图:下方的彩色方格图代表类别概率映射。每个网格单元预测了\( C \)个类别的概率,这个例子中有20个标记类别(\( C = 20 \))。
4.最终检测:经过处理后的图像展示了最终检测结果,其中不同颜色的边界框表示不同对象的位置,并附带置信度评分,说明了模型的预测。
5.编码为张量:所有这些预测——边界框、置信度和类别概率——都被编码为一个\( S \times S \times (B \times 5 + C) \)的张量。对于\( S = 7 \), \( B = 2 \), \( C = 20 \),最终的预测结果是一个\( 7 \times 7 \times 30 \)的张量。
1.统一的检测流程:
○YOLO将传统对象检测流程中的不同部分,如特征提取、边界框预测、类别概率估计等,融合到一个单一的神经网络中。这意味着整个检测流程从图像输入到边界框和类别的预测都在一个网络中完成。
2.图像分割与网格单元责任制:
○输入图像被划分为一个S×S的网格。如果一个对象的中心落在某个网格单元内,那么这个网格单元就负责检测该对象。
中心点定位:对于每个对象,YOLO定位其边界框(通常是由数据标注提供的真实边界框),并计算该边界框的中心点坐标
工作原理:
○输入图像首先被缩放到网络要求的固定尺寸(例如448x448像素)。
○缩放后的图像被划分为S×S的网格(如7×77×7)。
○如果对象的中心点落在一个特定的网格单元内,该网格单元就负责预测该对象的边界框和类别。
def find_object_centers(objects, grid_size):
centers = []
for obj in objects:
# 假设边界框的格式为 [x_min, y_min, x_max, y_max]
bbox = obj['bbox']
# 计算中心点坐标
center_x = (bbox[0] + bbox[2]) / 2
center_y = (bbox[1] + bbox[3]) / 2
# 判定中心点落在哪个网格单元
grid_x = int(center_x / grid_size)
grid_y = int(center_y / grid_size)
# 将中心点坐标和所属网格单元添加到列表
centers.append((center_x, center_y, grid_x, grid_y))
return centers
# 示例对象边界框数据
objects = [{'bbox': [50, 50, 150, 150]}, {'bbox': [300, 300, 400, 400]}]
# 网格尺寸
grid_size = 448 / 7
# 获取每个对象中心点及其所在的网格单元
object_centers = find_object_centers(objects, grid_size)
3.边界框和置信度预测:
○每个网格单元预测B个边界框及其置信度分数,这些置信度分数反映了模型对于边界框内存在对象的置信程度,以及对其预测位置的准确性。置信度被定义为对象存在的概率乘以预测框与真实框的交并比(IOU)。
工作原理:
▪每个边界框由5个预测组成:x,y,w,h和置信度conf。
•x,y是边界框中心相对于网格单元边界的位置。
•w,h是边界框的宽度和高度,相对于整个图像的尺寸。
•置信度conf是模型对边界框内包含对象的置信程度和预测框与真实框的交并比(IOU)的乘积。
▪每个网格单元还会预测属于每个类别的条件概率
def predict_confidence_and_bbox(grid_cell, objectness, bbox_params):
# objectness是模型预测的对象存在的概率
# bbox_params是模型预测的边界框参数,包含(x, y, w, h)
# 真实边界框信息通常在训练期间由数据集提供
# 真实边界框和预测边界框的IOU计算
iou = calculate_iou(bbox_params, ground_truth_bbox)
# 置信度分数计算
confidence = objectness * iou
return confidence, bbox_params
# 假设某个网格单元的对象存在概率和边界框参数
objectness = 0.8
bbox_params = {'x': 0.5, 'y': 0.5, 'w': 0.2, 'h': 0.2}
# 真实边界框信息,这里为了示例使用硬编码的值
ground_truth_bbox = {'x': 0.45, 'y': 0.55, 'w': 0.25, 'h': 0.25}
# 对于网格单元进行置信度和边界框预测
confidence, predicted_bbox = predict_confidence_and_bbox(
grid_cell='some_grid_cell',
objectness=objectness,
bbox_params=bbox_params
)
4.边界框细节:
○每个边界框的预测包括5个参数:中心坐标(x, y)、宽度w、高度h和置信度。中心坐标是相对于网格单元边界的,而宽度和高度是相对于整个图像的。
5.类别概率预测:
○每个网格单元还预测C个条件类别概率,这些概率基于网格单元内包含对象的假设。每个网格单元只预测一组类别概率,不管它预测了多少个边界框。
6.测试时的计算:
○在测试时,将条件类别概率与个别边界框的置信度预测相乘,得到每个边界框的类别特定置信度分数。这些分数既编码了特定类别出现在边界框中的概率,也编码了预测框与对象的拟合程度。
通过以上步骤,YOLO能够快速且准确地检测图像中的多个对象及其类别,实现了对图像内容的全局理解和实时处理。
模型:
在这里我们可以顺道复习一下卷积神经网络的知识:
这里以一个例子来解释参数的意思:
卷积层"Conv. Layer 7x7x64-s-2":表示卷积核大小为7*7,64个卷积核数(也称特征通道数),-2表示卷积核移动的步长为2
○这里可以重点理解一下64个卷积核工作原理,每一个卷积核都会单独的对原图像进行一遍卷积运算得到一个特征图像。64个卷积核分别运算后会得到64个特征图像,因此可以理解为原图像在64个维度下提取的特征图
池化层"Maxpool Layer 2x2-s-2":表示使用2x2的窗口进行最大池化,步长为2
全连接层
从7*7*1024到4096
▪ 最后一个卷积层输出的是一个具有7x7x1024尺寸的特征图。全连接层需要将这个三维的特征图展平(Flatten)成一个一维的向量,然后通过全连接操作转换为一个长度为4096的向量。展平操作本质上是将多维数组重塑为一维数组。在这种情况下,7x7x1024的特征图会被展平为一个具有 7×7×1024=501767×7×1024=50176个元素的一维数组。接着,这个一维数组会被送入全连接层。在全连接层中,每个输入节点会被连接到每个输出节点。在这里,如果输出层的大小是4096,那么会有一个50176×409650176×4096的权重矩阵。输入向量与权重矩阵相乘(还有偏置向量相加)得到了一个长度为4096的一维输出向量。
•注:50176×4096的权重矩阵的值是随机初始化的
○## 伪代码
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.models import Sequential
model = Sequential()
# 假设此前的层已经添加到模型中,最后的特征图尺寸为7x7x1024
# ...
# 添加Flatten层,将特征图展平成一维向量
model.add(Flatten()) # 这会将7x7x1024的特征图展平为一个50176元素的一维向量
# 添加全连接层,输出尺寸为4096
model.add(Dense(4096, activation='relu'))
# 可以继续添加更多的层
# 例如,如果需要输出30个值(例如YOLO的最后一层,通常包含边界框的坐标和类别的置信度)
model.add(Dense(30, activation='linear')) # 此处的激活函数取决于具体任务
# 接下来可以编译模型、训练或进行预测等
从4096到7*7*30
○先解释一下7*7*30的含义:
YOLO算法的最后一个层输出的是一个三维张量,其维度通常是 [S × S × (B * 5 + C)]。其中:
•S × S:表示将输入图像划分成了 S × S 个格子(cells)或网格(grid)。每个网格负责检测以它为中心的物体。
•B:表示每个网格预测的边界框(bounding box)数量。YOLO使用了多个边界框来捕捉不同形状的物体。
•5:表示每个边界框包含的信息量,这通常包括4个坐标值(中心点 x,中心点 y,宽度 w,高度 h)和1个置信度(confidence)得分,这个得分表示边界框中有物体的概率以及边界框的准确度。
•C:表示类别的数量,即模型能够检测的不同物体类别的数目。
每个边界框的四个坐标值经过了特殊处理:
•x 和 y:预测值是相对于各自网格单元的位置偏移,并通过sigmoid函数映射到0到1之间,表示在该网格单元内物体中心的相对位置。
•w 和 h:预测值是边界框宽度和高度的相对于整张图像的尺寸的对数空间预测,可能需要对其进行指数转换来获取实际的宽度和高度。
置信度得分反映了预测的边界框内存在物体的概率,并且这个得分也通过sigmoid函数映射到0到1之间。
每个类别的概率是在每个网格单元层面上进行预测的,表示在这个网格单元中存在每个类别物体的概率。
最终,这些值组合起来用于在测试时通过非最大抑制(Non-maximum Suppression, NMS)来确定最终的检测结果。NMS负责去除重叠边界框,保留最佳候选。
在这篇论文里:7*7(2*5+20)
○这里需要注意:C的值不是一张图片中识别的类别数量,而是在特定场景下定义的模型可以识别的总的类别数,与数据集和标注有关
"C" 在 YOLO 的最后输出层中指的是类别的数量,并不是指一张图片中能够检测的物体的数量,而是指模型被训练识别的物体类别的总数。换句话说,这个数字代表了模型可以区分的不同类别的种类数量。
举个例子,如果你有一个用于常见交通工具识别的数据集,它包含了汽车、摩托车、公交车、卡车、自行车五种交通工具的图像,那么在这种情况下 C 的值就是 5
这里回顾一下输出特征图(也称为卷积层输出)的尺寸大小计算公式:
为什么要进行预训练学习:
预训练模型通常在一个大型数据集上训练,比如ImageNet,这样可以学习到一组丰富的特征表示。这里提到的“预训练”,指的是在进行物体检测任务之前,先在图像分类任务上训练卷积层。具体到YOLO算法,预训练具有以下几个目的和好处:
1.知识转移:在ImageNet这样的大型数据集上预训练可以帮助网络学习到从基础特征(如边缘和纹理)到更高级的视觉概念(如物体的各个部分)的丰富表征。这些特征在许多视觉任务中是通用的,因此可以转移到新的任务上,比如物体检测。
2.收敛速度:预训练的网络通常能够更快地收敛,因为它已经学习了一些通用的特征,所以只需要调整这些特征使其适应新任务。
3.数据不足:对于某些特定的应用,可能没有足够的标注数据来从头训练一个复杂的深度网络。预训练可以在这种情况下提供帮助,因为预训练模型可以在有限的数据上进行微调(fine-tuning)。
4.提高性能:即使在数据量充足的情况下,预训练模型也通常会比从头开始训练的模型有更好的性能,因为它们可以利用在大规模数据集上学习到的知识。
损失函数
•
论文中Loss函数,密密麻麻的公式初看可能比较难懂。其实论文中给出了比较详细的解释。所有的损失都是使用平方和误差公式。
(1)预测框的中心点(x,y)。造成的损失是上图中的第一行。其中IIijobj为控制函数,在标签中包含物体的那些格点处,该值为 1 ;若格点不含有物体,该值为 0。也就是只对那些有真实物体所属的格点进行损失计算,若该格点不包含物体,那么预测数值不对损失函数造成影响。(x,y)数值与标签用简单的平方和误差。(2)预测框的宽高。造成的损失是上图的第二行。IIijobj的含义一样,也是使得只有真实物体所属的格点才会造成损失。这里对在损失函数中的处理分别取了根号,原因在于,如果不取根号,损失函数往往更倾向于调整尺寸比较大的预测框。例如,20个像素点的偏差,对于800乘600的预测框几乎没有影响,此时的IOU数值还是很大,但是对于30乘40的预测框影响就很大。取根号是为了尽可能的消除大尺寸框与小尺寸框之间的差异。(3)第三行与第四行,都是预测框的置信度C。当该格点不含有物体时,该置信度的标签为0;若含有物体时,该置信度的标签为预测框与真实物体框的IOU数值(IOU计算公式为:两个框交集的面积除以并集的面积)。
(4)第五行为物体类别概率P,对应的类别位置,该标签数值为1,其余位置为0,与分类网络相同。
此时再来看入coord与入noobj,Yolo面临的物体检测问题,是一个典型的类别数目不均衡的问题。其中49个格点,含有物体的格点往往只有3、4个,其余全是不含有物体的格点。此时如果不采取点措施,那么物体检测的mAP不会太高,因为模型更倾向于不含有物体的格点。入coord与入noobj的作用,就是让含有物体的格点,在损失函数中的权重更大,让模型更加“重视”含有物体的格点所造成的损失。在论文中, 取值分别为5与0.5。
源码解析:
如何从网络输出的7*7*30解析出每一个边界框的坐标,置信度以及类别的概率?
我在读这篇论文的时候,读到后面一直在思考一个问题,这个网络架构设计的确实巧妙,但是如何反解码出检测框的信息呢?
这里我们可以看看Yolo v1官方给的decode源码:
def decode(self, pred_tensor):
""" Decode tensor into box coordinates, class labels, and probs_detected.
Args:
pred_tensor: (tensor) tensor to decode sized [S, S, 5 x B + C], 5=(x, y, w, h, conf)
Returns:
boxes: (tensor) [[x1, y1, x2, y2]_obj1, ...]. Normalized from 0.0 to 1.0 w.r.t. image width/height, sized [n_boxes, 4].
labels: (tensor) class labels for each detected boxe, sized [n_boxes,].
confidences: (tensor) objectness confidences for each detected box, sized [n_boxes,].
class_scores: (tensor) scores for most likely class for each detected box, sized [n_boxes,].
"""
S, B, C = self.S, self.B, self.C
boxes, labels, confidences, class_scores = [], [], [], []
cell_size = 1.0 / float(S)
conf = pred_tensor[:, :, 4].unsqueeze(2) # [S, S, 1]
for b in range(1, B):
conf = torch.cat((conf, pred_tensor[:, :, 5*b + 4].unsqueeze(2)), 2)
conf_mask = conf > self.conf_thresh # [S, S, B]
# TBM, further optimization may be possible by replacing the following for-loops with tensor operations.
for i in range(S): # for x-dimension.
for j in range(S): # for y-dimension.
class_score, class_label = torch.max(pred_tensor[j, i, 5*B:], 0)
for b in range(B):
conf = pred_tensor[j, i, 5*b + 4]
prob = conf * class_score
if float(prob) < self.prob_thresh:
continue
# Compute box corner (x1, y1, x2, y2) from tensor.
box = pred_tensor[j, i, 5*b : 5*b + 4]
x0y0_normalized = torch.FloatTensor([i, j]) * cell_size # cell left-top corner. Normalized from 0.0 to 1.0 w.r.t. image width/height.
xy_normalized = box[:2] * cell_size + x0y0_normalized # box center. Normalized from 0.0 to 1.0 w.r.t. image width/height.
wh_normalized = box[2:] # Box width and height. Normalized from 0.0 to 1.0 w.r.t. image width/height.
box_xyxy = torch.FloatTensor(4) # [4,]
box_xyxy[:2] = xy_normalized - 0.5 * wh_normalized # left-top corner (x1, y1).
box_xyxy[2:] = xy_normalized + 0.5 * wh_normalized # right-bottom corner (x2, y2).
# Append result to the lists.
boxes.append(box_xyxy)
labels.append(class_label)
confidences.append(conf)
class_scores.append(class_score)
if len(boxes) > 0:
boxes = torch.stack(boxes, 0) # [n_boxes, 4]
labels = torch.stack(labels, 0) # [n_boxes, ]
confidences = torch.stack(confidences, 0) # [n_boxes, ]
class_scores = torch.stack(class_scores, 0) # [n_boxes, ]
else:
# If no box found, return empty tensors.
boxes = torch.FloatTensor(0, 4)
labels = torch.LongTensor(0)
confidences = torch.FloatTensor(0)
class_scores = torch.FloatTensor(0)
return boxes, labels, confidences, class_scores
这个源码里( conf = pred_tensor[:, :, 4].unsqueeze(2))结合论文里给的7*7*30的含义([S × S × (B * 5 + C)])
大概就能明白,其实就是一一映射的关系,7*7*30,我们可以这样解读一下:7*7个单元格,每个单元格有30个维度的特征,一个单元格负责两个检测框,每个检测框有5个信息(四个点坐标,置信度),所以30=5+5+20,前五个特征是第一个边界框的信息,第6-10个特质(第二组5个数据)是第二个边界框的信息,后面20个是类标签的概率。这样我们就能直接从输出的7*7*30的特征张量中解析出两个边界框的信息,以及20个类别的概率,类别概率在两个边界框中是共享的!
其他的代码就一目了然,坐标信息和置信度就是按照这个思路解析出来,最后将提取的B个检测框的张量堆叠一下就好。
NMS(non maximum supression非最大抑制的源码,用来消除冗余检测框的)
def nms(self, boxes, scores):
""" Apply non maximum supression.
Args:
Returns:
"""
threshold = self.nms_thresh
x1 = boxes[:, 0] # [n,]
y1 = boxes[:, 1] # [n,]
x2 = boxes[:, 2] # [n,]
y2 = boxes[:, 3] # [n,]
areas = (x2 - x1) * (y2 - y1) # [n,]
_, ids_sorted = scores.sort(0, descending=True) # [n,]
ids = []
while ids_sorted.numel() > 0:
# Assume `ids_sorted` size is [m,] in the beginning of this iter.
i = ids_sorted.item() if (ids_sorted.numel() == 1) else ids_sorted[0]
ids.append(i)
if ids_sorted.numel() == 1:
break # If only one box is left (i.e., no box to supress), break.
inter_x1 = x1[ids_sorted[1:]].clamp(min=x1[i]) # [m-1, ]
inter_y1 = y1[ids_sorted[1:]].clamp(min=y1[i]) # [m-1, ]
inter_x2 = x2[ids_sorted[1:]].clamp(max=x2[i]) # [m-1, ]
inter_y2 = y2[ids_sorted[1:]].clamp(max=y2[i]) # [m-1, ]
inter_w = (inter_x2 - inter_x1).clamp(min=0) # [m-1, ]
inter_h = (inter_y2 - inter_y1).clamp(min=0) # [m-1, ]
inters = inter_w * inter_h # intersections b/w/ box `i` and other boxes, sized [m-1, ].
unions = areas[i] + areas[ids_sorted[1:]] - inters # unions b/w/ box `i` and other boxes, sized [m-1, ].
ious = inters / unions # [m-1, ]
# Remove boxes whose IoU is higher than the threshold.
ids_keep = (ious <= threshold).nonzero().squeeze() # [m-1, ]. Because `nonzero()` adds extra dimension, squeeze it.
if ids_keep.numel() == 0:
break # If no box left, break.
ids_sorted = ids_sorted[ids_keep+1] # `+1` is needed because `ids_sorted[0] = i`.
return torch.LongTensor(ids)
while ids_sorted.numel() > 0:
# Assume ids_sorted size is [m,] in the beginning of this iter.
i = ids_sorted.item() if (ids_sorted.numel() == 1) else ids_sorted[0]
ids.append(i)
if ids_sorted.numel() == 1:
break # If only one box is left (i.e., no box to supress), break.
inter_x1 = x1[ids_sorted[1:]].clamp(min=x1[i]) # [m-1, ]
inter_y1 = y1[ids_sorted[1:]].clamp(min=y1[i]) # [m-1, ]
inter_x2 = x2[ids_sorted[1:]].clamp(max=x2[i]) # [m-1, ]
inter_y2 = y2[ids_sorted[1:]].clamp(max=y2[i]) # [m-1, ]
inter_w = (inter_x2 - inter_x1).clamp(min=0) # [m-1, ]
inter_h = (inter_y2 - inter_y1).clamp(min=0) # [m-1, ]
inters = inter_w * inter_h # intersections b/w/ box i and other boxes, sized [m-1, ].
unions = areas[i] + areas[ids_sorted[1:]] - inters # unions b/w/ box i and other boxes, sized [m-1, ].
ious = inters / unions # [m-1, ]
# Remove boxes whose IoU is higher than the threshold.
ids_keep = (ious <= threshold).nonzero().squeeze() # [m-1, ]. Because nonzero() adds extra dimension, squeeze it.
if ids_keep.numel() == 0:
break # If no box left, break.
ids_sorted = ids_sorted[ids_keep+1] # +1 is needed because ids_sorted[0] = i.
return torch.LongTensor(ids)
detect部分的代码(实际上就是把从yolo网络里跑出来特征decode一下,然后再做一些处理,最后绘制框框即可,思路很简答)
Python
def detect(self, image_bgr, image_size=448):
""" Detect objects from given image.
Args:
image_bgr: (numpy array) input image in BGR ids_sorted, sized [h, w, 3].
image_size: (int) image width and height to which input image is resized.
Returns:
boxes_detected: (list of tuple) box corner list like [((x1, y1), (x2, y2))_obj1, ...]. Re-scaled for original input image size.
class_names_detected: (list of str) list of class name for each detected boxe.
probs_detected: (list of float) list of probability(=confidence x class_score) for each detected box.
"""
h, w, _ = image_bgr.shape
img = cv2.resize(image_bgr, dsize=(image_size, image_size), interpolation=cv2.INTER_LINEAR)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # assuming the model is trained with RGB images.
img = (img - self.mean) / 255.0
img = self.to_tensor(img) # [image_size, image_size, 3] -> [3, image_size, image_size]
img = img[None, :, :, :] # [3, image_size, image_size] -> [1, 3, image_size, image_size]
img = Variable(img)
img = img.cuda()
with torch.no_grad():
pred_tensor = self.yolo(img)
pred_tensor = pred_tensor.cpu().data
pred_tensor = pred_tensor.squeeze(0) # squeeze batch dimension.
# Get detected boxes_detected, labels, confidences, class-scores.
boxes_normalized_all, class_labels_all, confidences_all, class_scores_all = self.decode(pred_tensor)
if boxes_normalized_all.size(0) == 0:
return [], [], [] # if no box found, return empty lists.
# Apply non maximum supression for boxes of each class.
boxes_normalized, class_labels, probs = [], [], []
for class_label in range(len(self.class_name_list)):
mask = (class_labels_all == class_label)
if torch.sum(mask) == 0:
continue # if no box found, skip that class.
boxes_normalized_masked = boxes_normalized_all[mask]
class_labels_maked = class_labels_all[mask]
confidences_masked = confidences_all[mask]
class_scores_masked = class_scores_all[mask]
ids = self.nms(boxes_normalized_masked, confidences_masked)
boxes_normalized.append(boxes_normalized_masked[ids])
class_labels.append(class_labels_maked[ids])
probs.append(confidences_masked[ids] * class_scores_masked[ids])
boxes_normalized = torch.cat(boxes_normalized, 0)
class_labels = torch.cat(class_labels, 0)
probs = torch.cat(probs, 0)
# Postprocess for box, labels, probs.
boxes_detected, class_names_detected, probs_detected = [], [], []
for b in range(boxes_normalized.size(0)):
box_normalized = boxes_normalized[b]
class_label = class_labels[b]
prob = probs[b]
x1, x2 = w * box_normalized[0], w * box_normalized[2] # unnormalize x with image width.
y1, y2 = h * box_normalized[1], h * box_normalized[3] # unnormalize y with image height.
boxes_detected.append(((x1, y1), (x2, y2)))
class_label = int(class_label) # convert from LongTensor to int.
class_name = self.class_name_list[class_label]
class_names_detected.append(class_name)
prob = float(prob) # convert from Tensor to float.
probs_detected.append(prob)
return boxes_detected, class_names_detected, probs_detected