原文为英文，进行了翻译和部分修改，原文地址
代码地址：github仓库、ACgit仓库
第4部分：目标得分阈值化和非最大值抑制
? 这是从头实现 YOLO v3 检测器教程的第 4 部分。在上一部分中，我们实现了网络的前向传递。在这部分中，我们通过目标置信度和非最大抑制来阈值化检测。
? 在前面的部分中，我们建立了一个模型，该模型在给定输入图像的情况下输出多个目标检测。准确地说，输出是一个形状为 B x 10647 x 85 的张量。B 是一批图像的数量，10647 是每个图像预测的边界框数量，85 是边界框属性的数量。
? 然而，如第 1 部分所述，我们必须将输出进行objectness score阈值处理和非极大值抑制，以获得"真正检测"的内容。为此将在文件 util.py 中创建一个名为 write_results 的函数
def write_result(prediction, confidence, num_classes, nms_conf = 0.4):
? 该函数将prediction、confidence（objectness score阈值）、num_classes（在本例子中为 80）和 nms_conf（NMS IoU 阈值）作为输入。
目标置信度的阈值化
? 我们的prediction张量包含有关 B x 10647 个边界框的信息。对于objectness score低于阈值的每个边界框，将其每个属性（代表边界框的一整行）的值设置为零。
conf_mask = (prediction[:, :, 4] > confidence).float().unsqueeze(2) # torch.Size([1, 10647, 1])
prediction = prediction * conf_mask # 将objectness score低于confidence阈值的行清零，其余行不变
进行非最大抑制
? 现在拥有的边界框属性有中心坐标以及边界框的高度和宽度。但是，使用每个框的一对对角的坐标更容易计算两个框的 IoU。因此，我们将框的 (center x, center y, height, width) 属性转换为 (左上角 x, 左上角 y, 右下角 x, 右下角 y)。
box_corner = prediction.new(prediction.shape)
# 左上角
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2 # 中心x-宽/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
? 如前所述，write标志用于指示尚未初始化输出，我们将使用一个张量来收集整个批次的"真正检测"。
? 一旦进入循环，请注意，每个边界框的那一行都有 85 个属性，其中 80 个是class scores。在这一点上，我们只关心具有最大值的class scores。因此，我们从每行中删除 80 个class scores，然后添加具有最大值的类别的索引以及该类别的class scores。
# confidence threshholding
# NMS
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
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
# confidence is at the top
conf_sort_index = torch.sort(image_pred_class[:, 4], descending=True)[1] # 获得objectness score降序排列的索引
image_pred_class = image_pred_class[conf_sort_index] # 对当前类别下的检测框按objectness score降序排列
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:
# 计算当前检测框和后面所有检测框的IOU，舍弃掉IOU大于阈值的检测框
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 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 zero entries
non_zero_ind = torch.nonzero(image_pred_class[:, 4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind].view(-1, 7)
? 另请注意，我们已将计算 iou 的代码行放在 try-catch 块中。这是因为循环旨在运行 idx 迭代（image_pred_class 中的行数）。然而，当我们继续循环时，可能会从 image_pred_class 中删除一些边界框。这意味着，即使从 image_pred_class 中删除一个值，我们也不能进行 idx 迭代。因此，我们可能会索引一个越界的值（IndexError），或者切片 image_pred_class[i+1:] 可能会返回一个空的张量，分配它会触发 ValueError。如果出现这两种错误，可以确定 NMS 不能移除更多的边界框，就跳出循环。
计算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
# 计算两个box左上角点坐标的最大值和右下角坐标的最小值
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个角坐标、objectness score、置信度最大的类的score和该类的索引。
? 和以前一样，除非我们有一个检测要分配，否则我们不会初始化我们的output张量。一旦输出张量被初始化，就将后续的检测连接到它。使用write标志来指示张量是否已初始化。在迭代类的循环结束时，将结果检测添加到张量output中。
batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
# Repeat the batch_ind for as many detections of the class cls in the image
seq = batch_ind, image_pred_class # shape （-1， 8）
if not write:
output = torch.cat(seq, 1)
write = True
else:
out = torch.cat(seq, 1)
output = torch.cat((output, out)) # 按行连接
? 在函数的最后，检查输出是否已经初始化。如果没有，则意味着在该批次的任何图像中都没有进行过一次检测。在这种情况下返回 0。
try:
return output # (D, 8)
except:
return 0
? 在这篇文章的最后，最终得到了一个张量形式的预测，它列出了每个预测的行。现在唯一剩下的就是创建一个输入管道来从磁盘读取图像、计算预测、在图像上绘制边界框，然后显示/写入这些图像。这就是将在下一部分中做的事情。