目标检测nms python代码_从零开始 PyTorch 项目：YOLO v3 目标检测实现（下）

前几日，机器之心编译介绍了《从零开始 PyTorch 项目：YOLO v3 目标检测实现》的前 3 部分，介绍了 YOLO 的工作原理、创建 YOLO 网络层级和实现网络的前向传播的方法。本文包含了该教程的后面两个部分，将介绍「置信度阈值设置和非极大值抑制」以及「设计输入和输出流程」的方法。总体而言，本教程的目的是使用 PyTorch 实现基于 YOLO v3 的目标检测器，后者是一种快速的目标检测算法。

本教程使用的代码需要运行在 Python 3.5 和 PyTorch 0.3 版本之上。你可以在以下链接中找到所有代码：https://github.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch

所需背景知识

1.本教程 1-3 部分

2.了解 PyTorch 基本工作方式，包括使用 nn.Module、nn.Sequential 和 torch.nn.parameter 类创建自定义架构的方式

3.NumPy 基本知识

4.OpenCV 基本知识

如果你缺少这些预备知识，可参阅文末扩展阅读部分了解。

置信度阈值设置和非极大值抑制

在前面 3 部分中，我们已经构建了一个能为给定输入图像输出多个目标检测结果的模型。具体来说，我们的输出是一个形状为 B x 10647 x 85 的张量；其中 B 是指一批(batch)中图像的数量，10647 是每个图像中所预测的边界框的数量，85 是指边界框属性的数量。

但是，正如第 1 部分所述，我们必须使我们的输出满足 objectness 分数阈值和非极大值抑制(NMS)，以得到后文所说的「真实(true)」检测结果。要做到这一点，我们将在 util.py 文件中创建一个名为 write_results 的函数。def write_results(prediction, confidence, num_classes, nms_conf = 0.4):

该函数的输入为预测结果、置信度(objectness 分数阈值)、num_classes(我们这里是 80)和 nms_conf(NMS IoU 阈值)。

目标置信度阈值

我们的预测张量包含有关 B x 10647 边界框的信息。对于有低于一个阈值的 objectness 分数的每个边界框，我们将其每个属性的值(表示该边界框的一整行)都设为零。conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)

prediction = prediction*conf_mask

执行非极大值抑制

注：我假设你已经理解 IoU(Intersection over union)和非极大值抑制(Non-maximum suppression)的含义了。如果你还不理解，请参阅文末提供的链接。

我们现在拥有的边界框属性是由中心坐标以及边界框的高度和宽度决定的。但是，使用每个框的两个对角坐标能更轻松地计算两个框的 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 的 batch 中有 1、2、3 这 3 张图像，它们各自有 5、2、4 个「真实」检测结果。因此，一次只能完成一张图像的置信度阈值设置和 NMS。也就是说，我们不能将所涉及的操作向量化，而且必须在预测的第一个维度(包含一个 batch 中图像的索引)上循环。batch_size = prediction.size(0)

write = False

for ind in range(batch_size):

image_pred = prediction[ind]          #image Tensor

#confidence threshholding

#NMS

如前所述，write 标签是用于指示我们尚未初始化输出，我们将使用一个张量来收集整个 batch 的「真实」检测结果。

进入循环后，我们再更清楚地说明一下。注意每个边界框行都有 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)

记得我们将 object 置信度小于阈值的边界框行设为零了吗？让我们摆脱它们。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 表示)的检测结果。

注意，以下代码在原始代码文件中有 3 格缩进，但因为页面空间有限，这里没有缩进。#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 = b