从头实现YOLOv3：第4部分-CSDN博客

本文链接：https://blog.csdn.net/ManiacLook/article/details/121503304

原文为英文，进行了翻译和部分修改，原文地址
代码地址：github仓库、ACgit仓库
相关内容：
YOLOv3论文翻译
 YOLOv3原理及流程简述
 从头实现YOLOv3：第1部分
 从头实现YOLOv3：第2部分
 从头实现YOLOv3：第3部分
 从头实现YOLOv3：第5部分

第4部分：目标得分阈值化和非极大值抑制

这是从头实现 YOLO v3 检测器教程的第 4 部分。在上一部分中，我们实现了网络的前向传递。在这部分中，我们通过目标置信度和非最大抑制来阈值化检测。

在前面的部分中，我们建立了一个模型，该模型在给定输入图像的情况下输出多个目标检测。准确地说，输出是一个形状为 B x 10647 x 85 的张量。B 是一批图像的数量，10647 是每个图像预测的边界框数量（3 x (13 x 13 + 26 x 26 + 52 x 52)），85 是边界框属性的数量（4 + 1 + 80）。

然而，如第 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  # 中心y-高/2
    # 右下角
    box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2  # 中心x+宽/2
    box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2  # 中心y+高/2

每个图像中"真正检测"的数量可能不同。例如，一个大小为 3 的批次，其中图像 1、2 和 3 分别有 5、2、4 个真实检测。因此，必须一次对一张图像进行置信度阈值化和 NMS。这意味着，我们不能矢量化所涉及的操作，并且必须循环prediction的第 0 维（批次中的图像索引）。

	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)  # 最大置信度 torch.Size([10647, 1])
        max_conf_score = max_conf_score.float().unsqueeze(1)  # 最大置信度索引 torch.Size([10647, 1])
        seq = (image_pred[:, :5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)

还记得我们将目标置信度小于阈值的边界框的那一行设置为零吗？直接去掉它们。

		# 将 objectness score 低于 confidence 阈值的行去掉
		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)
            # image_pred_class 只保留当前类的检测框，并且去除置信度为 0 的检测框
            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。

bbox-3

如果有两个相同类别的边界框的 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 中删除一些边界框。因此，我们可能会索引一个越界的值（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 中的图像索引
			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

在这篇文章的最后，最终得到了一个张量形式的预测，它列出了每个预测框的行。现在唯一剩下的就是创建一个输入管道来从磁盘读取图像、计算预测框、在图像上绘制边界框，然后显示/写入这些图像。这就是将在下一部分中做的事情。