YOLOV3Pytorch实战教程Part4：深入解析置信度阈值与非极大值抑制技术

前一节我们实现了网络的前向传播。这一节我们对检测输出设置目标置信度阈值和进行非极大值抑制。

必要条件&＃xff1a;

1.此系列教程的Part1到Part3。

2.Pytorch的基本知识&＃xff0c;包括如何使用nn.Module&＃xff0c;nn.Sequential&＃xff0c;torch.nn.parameter类构建常规的结构

3.numpy的基础知识

此前我们已经建立了一个模型&＃xff0c;给定一张输入图片它能产生B*10674*85维的输出向量。B是批中图片的数目&＃xff0c;10674是每张图片预测的边界框数目&＃xff0c;85是边界框属性数目。

但就像我们在part1中描述的那样&＃xff0c;我们必须对输出进行目标置信度阈值化和非极大值抑制&＃xff0c;以获得最终剩余的真正检测。为此&＃xff0c;我们将在文件util.py中创建一个名为write_results的函数。

def write_results(prediction, confidence, num_classes, nms_conf &＃61; 0.4):

这个函数将prediction、confidence(目标置信度阈值)、num_classes(在我们的示例中是80)和nms_conf (NMS IoU的阈值)作为输入。

目标置信度阈值&＃xff1a;

我们的预测张量包含了关于B x 10647个边界框的信息。对于每个目标置信度低于阈值的边界框&＃xff0c;我们将它的每个属性(行向量)的值设置为零。

conf_mask &＃61; (prediction[:,:,4] > confidence).float().unsqueeze(2)

prediction&＃61; prediction*conf_mask

执行极大值抑制&＃xff1a;

我们现在拥有的是边界框的中心坐标以及高度和宽度&＃xff0c;然而使用边界框的对角点更容易计算IOU。因此&＃xff0c;我们将框的(center x, center y, height, width)属性转换为(左上角x&＃xff0c;左上角y&＃xff0c;右下角x&＃xff0c;右下角y)。

box_corner &＃61;prediction.new(prediction.shape)

box_corner[:,:,0]&＃61; (prediction[:,:,0] - prediction[:,:,2]/2)

box_corner[:,:,1] &＃61; (prediction[:,:,1] - prediction[:,:,3]/2)

box_corner[:,:,2] &＃61; (prediction[:,:,0] &＃43; prediction[:,:,2]/2)

box_corner[:,:,3] &＃61; (prediction[:,:,1] &＃43; prediction[:,:,3]/2)

prediction[:,:,:4] &＃61; box_corner[:,:,:4]

每幅图像中真实检测框的数目可能不同。例如&＃xff0c;一批大小为3的图像&＃xff0c;其中图像1、2和3分别有5、2、4个真检测值。因此&＃xff0c;每次必须对同一个图像进行置信阈值和NMS&＃xff0c;而不能对所涉及的操作进行矢量化&＃xff0c;必须在预测的第一个维度(包含成批图像的索引)上进行遍历操作。

batch_size &＃61;prediction.size(0)

write&＃61;Falsefor ind inrange(batch_size):

image_pred&＃61; prediction[ind] #image Tensor

#confidence threshholding

#NMS

write标志位用于指示我们是否对output进行了初始化&＃xff0c;将会使用一个向量来收集整个批中真实的预测。

循环的开始我们进行数据清理。因为每个边界框行有85个属性&＃xff0c;其中80个是类得分。我们只关心类得分最大值的那个&＃xff0c;所以会从每行中删除80个类得分&＃xff0c;添加具有最大值的类的索引&＃xff0c;以及该类的类得分。

max_conf, max_conf_score &＃61; torch.max(image_pred[:,5:5&＃43; num_classes], 1)

max_conf&＃61; max_conf.float().unsqueeze(1)

max_conf_score&＃61; max_conf_score.float().unsqueeze(1)

seq&＃61; (image_pred[:,:5], max_conf, max_conf_score)

image_pred&＃61; torch.cat(seq, 1)

我们前面已经将目标置信度得分低于阈值的边界框行属性设置为了0&＃xff0c;现在就筛除它们。

non_zero_ind &＃61; (torch.nonzero(image_pred[:,4]))try:

image_pred_&＃61; 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] &＃61;&＃61;0:continue

try-except块用于处理没有检测到的情况。在这种情况下&＃xff0c;我们使用continue跳过此图像的其余循环体。

接下来&＃xff0c;让我们在图像中检测目标。

#Get the various classes detected in the image

img_classes &＃61; unique(image_pred_[:,-1]) #-1 index holds the class index

因为对于同一个类别可能存在多个正确检测&＃xff0c;我们使用一个叫unique的函数来获得给定图片中所有出现的类。

defunique(tensor):

tensor_np&＃61;tensor.cpu().numpy()

unique_np&＃61;np.unique(tensor_np)

unique_tensor&＃61;torch.from_numpy(unique_np)

tensor_res&＃61;tensor.new(unique_tensor.shape)

tensor_res.copy_(unique_tensor)return tensor_res

之后我们对于每个类进行NMS

for cls inimg_classes:#perform NMS

一进入这个循环&＃xff0c;首先要做的事情就是提取对于某一特定类别的检测(用变量cls表示)

#get the detections with one particular class

cls_mask &＃61; image_pred_*(image_pred_[:,-1] &＃61;&＃61; cls).float().unsqueeze(1)

class_mask_ind&＃61; torch.nonzero(cls_mask[:,-2]).squeeze()

image_pred_class&＃61; image_pred_[class_mask_ind].view(-1,7)#sort the detections such that the entry with the maximum objectness#confidence is at the top

conf_sort_index &＃61; torch.sort(image_pred_class[:,4], descending &＃61; True )[1]

image_pred_class&＃61;image_pred_class[conf_sort_index]

idx&＃61; image_pred_class.size(0) #Number of detections

然后我们进行NMS

for i inrange(idx):#Get the IOUs of all boxes that come after the one we are looking at

#in the loop

try:

ious&＃61; bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i&＃43;1:])exceptValueError:break

exceptIndexError:break

#Zero out all the detections that have IoU > treshhold

iou_mask &＃61; (ious

image_pred_class[i&＃43;1:] *&＃61;iou_mask#Remove the non-zero entries

non_zero_ind &＃61; torch.nonzero(image_pred_class[:,4]).squeeze()

image_pred_class&＃61; image_pred_class[non_zero_ind].view(-1,7)

这里我们使用到了一个函数bbox_iou。第一个输入参数是循环体变量i索引处的边界框&＃xff0c;第二个输入参数是多行边界框的一个tensor。函数bbox_iou的输出是一个tensor它包含了第一个输入的边界框与第二个输入的所有边界框的IOU。如下&＃xff1a;

之前我们已经将目标置信度高的边界框放在前面&＃xff0c;如果后面的边界框IoU值与前面的相比超过了阈值&＃xff0c;那后者就会被删去。

循环体里面下面这行计算IoU。

ious &＃61; bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i&＃43;1:])

每轮迭代&＃xff0c;如果有任何索引大于i的边界框与第i个边界框的IoU大于阈值nms_thresh&＃xff0c;那这个边界框就会被删除。

#Zero out all the detections that have IoU > treshhold

iou_mask &＃61; (ious

image_pred_class[i&＃43;1:] *&＃61;iou_mask#Remove the non-zero entries

non_zero_ind &＃61; torch.nonzero(image_pred_class[:,4]).squeeze()

image_pred_class&＃61; image_pred_class[non_zero_ind]

还要注意的是&＃xff0c;我们将计算IoU的代码行放在try-catch块中。这是因为此循环按照id进行迭代(image_pred_class中的行数)。但因为我们循环过程中可能会从image_pred_class中删除一些边界框。这样一来&＃xff0c;迭代可能会出现索引越界触发IndexError或者image_pred_class[i&＃43;1:]返回一个空张量触发ValueError。此时我们可以确定NMS已经无法删除多余的边界框了&＃xff0c;从而跳出循环。

计算IoU:

defbbox_iou(box1, box2):"""Returns the IoU of two bounding boxes"""

#Get the coordinates of bounding boxes

b1_x1, b1_y1, b1_x2, b1_y2 &＃61; box1[:,0], box1[:,1], box1[:,2], box1[:,3]

b2_x1, b2_y1, b2_x2, b2_y2&＃61; box2[:,0], box2[:,1], box2[:,2], box2[:,3]#get the corrdinates of the intersection rectangle

inter_rect_x1 &＃61;torch.max(b1_x1, b2_x1)

inter_rect_y1&＃61;torch.max(b1_y1, b2_y1)

inter_rect_x2&＃61;torch.min(b1_x2, b2_x2)

inter_rect_y2&＃61;torch.min(b1_y2, b2_y2)#Intersection area

inter_area &＃61; torch.clamp(inter_rect_x2 - inter_rect_x1 &＃43; 1, min&＃61;0) * torch.clamp(inter_rect_y2 - inter_rect_y1 &＃43; 1, min&＃61;0)#Union Area

b1_area &＃61; (b1_x2 - b1_x1 &＃43; 1)*(b1_y2 - b1_y1 &＃43; 1)

b2_area&＃61; (b2_x2 - b2_x1 &＃43; 1)*(b2_y2 - b2_y1 &＃43; 1)

iou&＃61; inter_area / (b1_area &＃43; b2_area -inter_area)return iou

写入预测&＃xff1a;

write_results函数输出一个形状为 Dx8 的tensor。这里D是所有图像的真实检测&＃xff0c;每个都用一行表示。每个检测有8个属性&＃xff0c;即检测所属批次图像的索引、4个角坐标、目标置信度得分、最大置信类得分、该类的索引。

和此前一样&＃xff0c;我们等到有一个检测时才初始化输出向量并将后续的检测拼接进来。使用写标志来表示tensor是否已经初始化。在遍历类的循环结束时&＃xff0c;我们将检测结果添加到输出tensor中。

batch_ind &＃61; 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 &＃61;batch_ind, image_pred_classif notwrite:

output&＃61; torch.cat(seq,1)

write&＃61;Trueelse:

out&＃61; torch.cat(seq,1)

output&＃61; torch.cat((output,out))

在函数的末尾&＃xff0c;我们检查输出是否已经初始化。如果没有&＃xff0c;就意味着这批图像中没有一个检测到。在这种情况下&＃xff0c;我们返回0。

try:returnoutputexcept:

return 0

这就是这一部分所要讲解的内容了。现在我们终于有了一个预测&＃xff0c;它以tensor的形式列出了每一个边界框。所以只剩下一件事就是创建一个输入管道来从磁盘读取图像&＃xff0c;计算预测&＃xff0c;在图像上绘制边界框&＃xff0c;然后显示/写入这些图像。这是我们下一部分要做的。

Further Reading