从零开始Pytorch-YOLOv3【笔记】（四）置信度阈值和非极大值抑制

前言

上一篇：从零开始Pytorch-YOLOv3【笔记】（三）实现网络的前向传播

上一篇我们实现了根据预训练权重通过前向网络传播输出了一个torch.Size([1, 10647, 85])的张量，其中 B=1 是指一批（batch）中图像的数量，10647 是每个图像中所预测的边界框的数量，85 是指边界框属性的数量（x,y,w,h,conf,cls）conf置信度，cls=80的COCO数据集。

对应从零开始 PyTorch 项目：YOLO v3 目标检测实现（下）

但是，正如第 1 部分所述，我们必须使我们的输出满足 objectness 分数阈值和非极大值抑制（NMS），以得到后文所说的「真实（true）」检测结果。

在这之前，我们先了解一下什么是目标置信度阈值和非极大值抑制

预备知识

目标置信度阈值

置信度：作者对他的作用定义有两重：
一重是：代表当前box是否有对象的概率\(P_{r}(Object)\)，注意，是对象，不是某个类别的对象，也就是说它用来说明当前box内只是个背景（backgroud）还是有某个物体（对象）；
另一重：表示当前的box有对象时，它自己预测的box与物体真实的box可能的\(IOU_{pred}^{truth}\) 的值，注意，这里所说的物体真实的box实际是不存在的，这只是模型表达自己框出了物体的自信程度。

经过以上的解释，其实我们也就可以用数学形式表示置信度的定义了：

\(C_{i}^{j} = P_{r}(Object) * IOU_{pred}^{truth}\)

其中，\(C_{i}^{j}\)表示第\(i\)个grid cell的第\(j\)个bounding box的置信度。
​
通过设置置信度阈值过滤置信度低的Bounding box可以避免很多的噪声。

非极大值抑制

该算法逐个去除冗余的Bounding box。它通过删除重叠大于我们设置的与ground truth的阈值来实现这一点。也就是IOU阈值要小于overThreshold

下面是NMS的全功能封装代码，接下来会详细解释它的原理。

点击查看代码

def NMS(boxes, overlapThresh = 0.4):
    # Return an empty list, if no boxes given
    if len(boxes) == 0:
        return []
    x1 = boxes[:, 0]  # x coordinate of the top-left corner
    y1 = boxes[:, 1]  # y coordinate of the top-left corner
    x2 = boxes[:, 2]  # x coordinate of the bottom-right corner
    y2 = boxes[:, 3]  # y coordinate of the bottom-right corner
    # Compute the area of the bounding boxes and sort the bounding
    # Boxes by the bottom-right y-coordinate of the bounding box
    areas = (x2 - x1 + 1) * (y2 - y1 + 1) # We add 1, because the pixel at the start as well as at the end counts
    # The indices of all boxes at start. We will redundant indices one by one.
    indices = np.arange(len(x1))
    for i,box in enumerate(boxes):
        # Create temporary indices  
        temp_indices = indices[indices!=i]
        # Find out the coordinates of the intersection box
        xx1 = np.maximum(box[0], boxes[temp_indices,0])
        yy1 = np.maximum(box[1], boxes[temp_indices,1])
        xx2 = np.minimum(box[2], boxes[temp_indices,2])
        yy2 = np.minimum(box[3], boxes[temp_indices,3])
        # Find out the width and the height of the intersection box
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)
        # compute the ratio of overlap
        overlap = (w * h) / areas[temp_indices]
        # if the actual boungding box has an overlap bigger than treshold with any other box, remove it's index  
        if np.any(overlap) > treshold:
            indices = indices[indices != i]
    #return only the boxes at the remaining indices
    return boxes[indices].astype(int)

该函数的第一个参数boxes是一个数组，每行代表一个Bounding box

第二个参数overThreshold限制了两个bounding box允许有的重叠区域。如果它们重叠得更多，那么其中一个就会被丢弃。overThreshold=0.4意味着两个bounding box最多可重叠40%的面积。

代码中我们通过找两个Bounding box的左上角坐标的最大值和右下角坐标的最小值找到重叠区域的两点坐标，也就得到了重叠区域的宽高。

然后计算重叠面积大于0.4的索引就会被删除。最后带有未删除索引的框。像素坐标必须是整数，所以为了安全起见，我们把它们转换成整数。

YOLOv3中的置信度阈值和非极大值抑制代码实现

该部分的全部代码如下：

点击查看代码

def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
    # 置信度过滤：将低于该置信度的Bounding box的Attribute全部置为零也就是最后一维的85个参数置为0
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)  # torch.Size([1, 10647, 1, 85])
    prediction = prediction*conf_mask
    # x_c, y_c, w, h -> x1, y1, x2, y2
    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]
    
    batch_size = prediction.size(0)

    write = False
    


    for ind in range(batch_size):
        # 对batch_size中的图像单独进行NMS
        image_pred = prediction[ind]          #image Tensor
        #confidence threshholding 
        #NMS
    
        max_conf, max_conf_score = torch.max(image_pred[:,5:5+ num_classes], 1)  # 第二个参数为1表示取每行的最大值，返回最大值value,索引index
        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]))  # 非0元素的索引
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        except:
            continue
        
        if image_pred_.shape[0] == 0:
            continue              

        #Get the various classes detected in the image
        img_classes = unique(image_pred_[:,-1])  # -1 index holds the class index
        
        
        for cls in img_classes:
            #perform NMS

        
            #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]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)   #Number of detections
            
            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)
                
            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))

    try:
        return output
    except:
        return 0

我们将在 util.py 文件中创建一个名为 write_results 的函数。

def write_results(prediction, confidence, num_classes, nms_conf = 0.4):

该函数的输入为预测结果、置信度（objectness 分数阈值）、num_classes（COCO数据集是80）和 nms_conf（NMS IoU 阈值）。

目标置信度阈值

> TODO: 不是很理解 这个张量的乘法做了什么

我们的预测张量包含有关 B x 10647 边界框的信息。对于有低于一个阈值的 objectness 分数的每个边界框，我们将其每个属性的值（表示该边界框的一整行）都设为零。

后续会通过torch.nonzero()对这一部分为0的b-box进行过滤。

    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction*conf_mask

详细解读一下这两行代码的作用：

(prediction[:,:,4] > confidence)这个代码得到predictiontorch.Size([1, 10647, 85])的第三维向量的85个参数中，第五列的置信度的值大于confidence则为True。结果格式为：

[[[False],
[False],
...,
[False],
...,
[True],
[True],
...,
[True]]]

(prediction[:,:,4] > confidence).float()之后是如下格式：

[[[0.],
[0.],
...,
[0.],
...,
[1.],
[1.],
...,
[1.]]]

(prediction[:,:,4] > confidence).float().unsqueeze(2)将torch.Size([1, 10647, 85])扩充为torch.Size([1, 10647, 1, 85]),结果如下：

[[[[0.],
[0.],
...,
[0.],
...,
[1.],
[1.],
...,
[1.]]]]

然后将10647个anchor通过张量乘法的方式按照上式如果conf_mask对应的行为0，则同样索引的anchor的attribute全为0。这样说有点抽象，看戏下面的例子就比较好理解了。(更多有关张量乘法的例子见：pytorch——张量乘法)

import numpy as np
import torch

arr = np.array([
    [11, 21, 31, 41, 51, 61, 71],
    [12, 22, 32, 42, 52, 62, 72],
    [13, 23, 33, 43, 53, 63, 73]
    ])
tensor = torch.tensor(arr)
print('tensor: \n', tensor)
print(tensor.size())

conf_mask = (tensor[:,4] > 52).float().unsqueeze(1)
prediction = tensor*conf_mask

print('conf_mask: \n', conf_mask)
print('----')
print('prediction: \n', prediction)

打印结果为：

tensor: 
 tensor([[11, 21, 31, 41, 51, 61, 71],
        [12, 22, 32, 42, 52, 62, 72],
        [13, 23, 33, 43, 53, 63, 73]], dtype=torch.int32)
torch.Size([3, 7])
conf_mask: 
 tensor([[0.],
        [0.],
        [1.]])
----
prediction: 
 tensor([[ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  0.,  0.,  0.],
        [13., 23., 33., 43., 53., 63., 73.]])

张量的乘法可以如图这样理解:

非极大值抑制

我们现在拥有的边界框属性是由中心坐标以及边界框的高度和宽度决定的。但是，使用每个框的两个对角坐标能更轻松地计算两个框的 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)  # 第二个参数为1表示取每行的最大值，返回最大值value,索引index
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)  # (x_c, y_c, w, h, conf, max_conf, max_conf_score)

然后使用torch.nonzero函数过滤掉先前我们将anchor的85Attribute全置为0的anchor。

non_zero_ind =  (torch.nonzero(image_pred[:,4]))  # max_conf这一列非0元素的索引
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  (max_conf_score这一列)

因为同一类别可能会有多个「真实」检测结果，所以我们使用一个名叫 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

这是对numpy.unique()进行了重构，返回torch的Tensor类型。

然后，我们按照类别执行 NMS。

for cls in img_classes:
    #perform NMS  # 按类别进行NMS

一旦我们进入循环，我们要做的第一件事就是提取特定类别（用变量 cls 表示）的检测结果。

    #get the detections with one particular class  # 只保留检测到当前cls类别的bbox
    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]  # 默认dim=-1，按行排序；descending = True递减排序
    image_pred_class = image_pred_class[conf_sort_index]  # 按照conf递减排序后的Tensor
    idx = image_pred_class.size(0)   #Number of detections # image_pred_class的行数

现在，我们执行 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  # 移除IOU < IOU阈值的bboxTensor
        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。

如果我们有 2 个同样类别的边界框且它们的 IoU 大于一个阈值，那么就去掉其中类别置信度较低的那个。我们已经对边界框进行了排序，其中有更高置信度的在上面。

在循环部分，bbox_iou给出了当前类别conf最高conf与它同一类别的 IoU（前面已经根据conf降序进行了排序）每次迭代时，如果有bbox的索引大于 i 且又大于阈值 nms_thresh 的 IoU（与索引为 i 的框），那么就去掉那个特定的bbox。

还要注意，我们已经将用于计算 ious 的代码放在了一个 try-catch 模块中。这是因为这个循环在设计上是为了运行 idx 次迭代（image_pred_class 中的行数）。但是，当我们继续循环时，一些边界框可能会从 image_pred_class 移除。这意味着，即使只从 image_pred_class 中移除了一个值，我们也不能有 idx 次迭代。因此，我们可能会尝试索引一个边界之外的值（IndexError），片状的 image_pred_class[i+1:] 可能会返回一个空张量，从而指定触发 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 = (inter_rect_x2 - inter_rect_x1 + 1)*(inter_rect_y2 - inter_rect_y1 + 1)

    #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 函数输出一个形状为 Dx8 的张量；其中 D 是所有图像中的「真实」检测结果，每个都用一行表示。每一个检测结果都有 8 个属性，即：该检测结果所属的 batch 中图像的索引、4 个角的坐标、objectness 分数、有最大置信度的类别的分数、该类别的索引。

如之前一样，我们没有初始化我们的输出张量，除非我们有要分配给它的检测结果。一旦其被初始化，我们就将后续的检测结果与它连接起来。我们使用 write 标签来表示张量是否初始化了。在类别上迭代的循环结束时，我们将所得到的检测结果image_pred_class加入到张量输出中。

    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  # 添加D行1列的batch_ind
    seq = batch_ind, image_pred_class

    if not write:  # 如果是当前batch的第一个，
        output = torch.cat(seq,1)
        write = True
    else:  # 如果是当前batch的后续，与原batch中的数据进行拼接
        out = torch.cat(seq,1)
        output = torch.cat((output,out))

在该函数结束时，我们会检查输出是否已被初始化。如果没有，就意味着在该 batch 的任意图像中都没有单个检测结果。在这种情况下，我们返回 0。

try:
    return output
except:
    return 0

这部分就到此为止了。在这部分结束时，我们终于有了一个张量形式的预测结果，其中以行的形式列出了每个预测。现在还剩下：创造一个从磁盘读取图像的输入流程，计算预测结果，在图像上绘制边界框，然后展示 / 写入这些图像。这是下一部分要介绍的内容。