如何在PyTorch中从头开始实现YOLO(v3)对象检测器:第4部分
这是从头开始实现YOLO v3探测器的教程的第4部分。在最后一部分,我们实施了网络的前向传递。在这一部分中,我们通过对象置信度和非最大抑制来阈值检测。
本教程的代码旨在在Python 3.5和PyTorch 0.4上运行。它可以在这个Github回购中找到它的全部内容。
先决条件
- 本教程的第1-3部分。
- PyTorch的基本工作知识,包括如何使用nn.Module,nn.Sequential和torch.nn.parameter类创建自定义体系结构。
- NumPy的基础知识
如果您缺少任何正面,帖子下面有链接供您关注。
在前面的部分中,我们构建了一个模型,该模型在给定输入图像的情况下输出多个对象检测。确切地说,我们的输出是一个形状的张量B x 10647 x 85
。B是批处理中的图像数,10647是每个图像预测的边界框数,85是边界框属性的数量。
但是,如第1部分所述,我们必须使输出符合对象分数阈值和非最大抑制,以获得我将在本文的其余部分中称为真实检测的内容。为此,我们将创建一个write_results
在文件中调用的函数util.py
def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
这些函数作为输入prediction
,confidence
(对象性得分阈值),num_classes
(在我们的例子中为80)和nms_conf
(NMS IoU阈值)。
对象置信度阈值
我们的预测张量包含有关B x 10647
边界框的信息。对于具有低于阈值的对象性得分的每个边界框,我们将其每个属性(表示边界框的整行)的值设置为零。
conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
prediction = prediction*conf_mask
执行非最大抑制
注意:我假设您了解IoU(联合交叉)是什么,以及非最大抑制是什么。如果不是这种情况,请参阅帖子末尾的链接)。
我们现在拥有的边界框属性由中心坐标以及边界框的高度和宽度来描述。但是,使用每个盒子的一对诊断角的坐标来计算两个盒子的IoU更容易。因此,我们将框的(中心x,中心y,高度,宽度)属性转换为(左上角x,左上角y,右下角x,右下角y)。
box_corner = prediction.new(prediction.shape)
box_corner[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
box_corner[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
box_corner[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2)
box_corner[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
prediction[:,:,:4] = box_corner[:,:,:4]
每个图像中的真实检测数可能不同。例如,批量为3的批次,其中图像1,2和3分别具有5,2,4个真实检测。因此,必须一次对一个图像进行置信度阈值处理和NMS。这意味着,我们无法对所涉及的操作进行矢量化,并且必须遍历第一维prediction
(包含批处理中的图像索引)。
batch_size = prediction.size(0)
write = False
for ind in range(batch_size):
image_pred = prediction[ind] #image Tensor
#confidence threshholding
#NMS
如前所述,write
flag用于表示我们尚未初始化output
,我们将使用一个张量来收集整个批次中的真实检测。
一旦进入循环,让我们清理一下。请注意,每个边界框行有85个属性,其中80个是类别分数。此时,我们只关注具有最大值的班级分数。因此,我们从每一行中删除80个类分数,而是添加具有最大值的类的索引,以及该类的类分数。
max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)
max_conf = max_conf.float().unsqueeze(1)
max_conf_score = max_conf_score.float().unsqueeze(1)
seq = (image_pred[:,:5], max_conf, max_conf_score)
image_pred = torch.cat(seq, 1)
还记得我们已经将具有小于阈值的对象置信度的边界框行设置为零吗?让我们摆脱他们。
non_zero_ind = (torch.nonzero(image_pred[:,4]))
try:
image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
except:
continue
#For PyTorch 0.4 compatibility
#Since the above code with not raise exception for no detection
#as scalars are supported in PyTorch 0.4
if image_pred_.shape[0] == 0:
continue
try-except块用于处理我们没有检测到的情况。在这种情况下,我们使用continue
跳过此图像的循环体的其余部分。
现在,让我们在图像中检测出类。
#Get the various classes detected in the image
img_classes = unique(image_pred_[:,-1]) # -1 index holds the class index
由于可以对同一个类进行多次真正的检测,因此我们使用一个函数调用unique
来获取任何给定图像中的类。
def unique(tensor):
tensor_np = tensor.cpu().numpy()
unique_np = np.unique(tensor_np)
unique_tensor = torch.from_numpy(unique_np)
tensor_res = tensor.new(unique_tensor.shape)
tensor_res.copy_(unique_tensor)
return tensor_res
然后,我们按类别执行NMS。
for cls in img_classes:
#perform NMS
一旦我们进入循环,我们要做的第一件事是提取特定类的检测(由变量表示cls
)。
以下代码在原始代码文件中缩进了三个块,但我没有在此处缩进,因为此页面上的空间有限。
#get the detections with one particular class
cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
image_pred_class = image_pred_[class_mask_ind].view(-1,7)
#sort the detections such that the entry with the maximum objectness
s#confidence is at the top
conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
image_pred_class = image_pred_class[conf_sort_index]
idx = image_pred_class.size(0) #Number of detections
现在,我们执行NMS。
for i in range(idx):
#Get the IOUs of all boxes that come after the one we are looking at
#in the loop
try:
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
except ValueError:
break
except IndexError:
break
#Zero out all the detections that have IoU > treshhold
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#Remove the non-zero entries
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind].view(-1,7)
在这里,我们使用一个函数bbox_iou
。第一个输入是由i
循环中的变量索引的边界框行。
第二个输入bbox_iou
是多行边界框的张量。函数的输出bbox_iou
是包含由第一输入表示的边界框的IoU的张量,其中每个边界框存在于第二输入中。
如果我们有两个具有大于阈值的IoU的同一类的边界框,则消除具有较低类置信度的边界框。我们已经整理出了具有更高置信度的边界框。
在循环体中,以下行给出了框的IoU,索引为i
所有具有高于的索引的边界框i
。
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
每次迭代,如果任何具有索引大于具有大于阈值i
的IoU(具有索引框i
)的边界框,nms_thresh
则消除该特定框。
#Zero out all the detections that have IoU > treshhold
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#Remove the non-zero entries
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind]
另请注意,我们已经ious
在try-catch块中放置了代码行来计算。这是因为循环被设计为运行idx
迭代(行数image_pred_class
)。但是,当我们继续循环时,可以从中删除许多边界框image_pred_class
。这意味着,即使删除了一个值image_pred_class
,我们也无法进行idx
迭代。因此,我们可能会尝试索引超出bounds(IndexError
)的值,或者切片image_pred_class[i+1:]
可能返回一个空张量,指定触发a的值ValueError
。此时,我们可以确定NMS不能删除任何进一步的边界框,并且我们打破了循环。
计算IoU
这是功能bbox_iou
。
def bbox_iou(box1, box2):
"""
Returns the IoU of two bounding boxes
"""
#Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:,0], box1[:,1], box1[:,2], box1[:,3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:,0], box2[:,1], box2[:,2], box2[:,3]
#get the corrdinates of the intersection rectangle
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
#Intersection area
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
#Union Area
b1_area = (b1_x2 - b1_x1 + 1)*(b1_y2 - b1_y1 + 1)
b2_area = (b2_x2 - b2_x1 + 1)*(b2_y2 - b2_y1 + 1)
iou = inter_area / (b1_area + b2_area - inter_area)
return iou
写预测
该函数write_results
输出形状D x 8的张量。这里D是所有图像中的真实检测,每个图像由一行表示。每个检测具有8个属性,即检测所属批次中的图像的索引,4个角坐标,对象得分,具有最大置信度的类的得分,以及该类的索引。
和以前一样,我们不会初始化输出张量,除非我们有一个检测分配给它。一旦初始化,我们将后续检测连接到它。我们使用该write
标志来指示张量是否已经初始化。在遍历类的循环结束时,我们将结果检测添加到张量output
。
batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
#Repeat the batch_id for as many detections of the class cls in the image
seq = batch_ind, image_pred_class
if not write:
output = torch.cat(seq,1)
write = True
else:
out = torch.cat(seq,1)
output = torch.cat((output,out))
在函数结束时,我们检查是否output
已经初始化。如果没有手段,则批次的任何图像中都没有单一检测。在那种情况下,我们返回0。
try:
return output
except:
return 0
这是这篇文章的内容。在这篇文章的最后,我们最终得到了一个张量形式的预测,它将每个预测列为行。现在唯一剩下的就是创建一个输入管道来从磁盘读取图像,计算预测,在图像上绘制边界框,然后显示/写入这些图像。这是我们将在下一部分中做的。