從零開始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

這部分就到此為止了。在這部分結束時，我們終於有了一個張量形式的預測結果，其中以行的形式列出了每個預測。現在還剩下：創造一個從磁盤讀取圖像的輸入流程，計算預測結果，在圖像上繪制邊界框，然后展示 / 寫入這些圖像。這是下一部分要介紹的內容。