part 4：置信度阈值化和非极大值抑制

翻译原文：https://blog.paperspace.com/how-to-implement-a-yolo-v3-object-detector-from-scratch-in-pytorch-part-3/

本篇文章是《如何使用PyTorch从零开始实现YOLO(v3)目标检测算法》的第四部分。这系列论文一共有五篇文章，五篇文章的主要内容在下文中有涉及。如果有问题欢迎和我交流~

上篇博客中，我们构造了一个网络模型：当给定一个输入图片的时候，能够输出一些目标对象。准确来说，我们的输出是一个B*10647*85的张量。B是一个批量中输入图片的数量，10647是每个图片预测边界框的数量，85是每个边界框的属性个数。我们必须对输出进行目标物体阈值化和非极大值抑制，来获得正确的预测。为了实现这个，我们将会在util.py文件创建一个名叫write_results的函数。

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

这个函数的输入是：prediction,confidence（目标物体阈值）,num_classes（在我们的例子中是80）,和nms_conf（NMS的IoU阈值）

1 目标物体得分阈值化

我们的预测包含B*10647个边界框的信息。对于目标物体得分低于阈值的边界框，我们把他的每一个属性的值设置为零（每一行代表一个边界框）

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

2 进行非极大值抑制

边界框的属性由边界框的中心坐标以及宽和高来表示。但是，使用每一个边界框的对角线坐标很容易计算两个边界框的IoU，因此，我们把我们边界框的属性值（中心x,中心y,w,h）变成(左上角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]

每张图片检测到目标物体的数量是各不相同的。举例来说，一个批量有三张图片1,2,3它们分别有5,2,4个目标物体。因此需要立即对每一张图片进行目标物体阈值化和非极大值抑制。我们必须循环prediction的第一个维度（包含一个批量的图片数量的索引）

    batch_size = prediction.size(0)write = Falsefor ind in range(batch_size):image_pred = prediction[ind]          #image Tensor#confidence threshholding #NMS

像之前描述的那样，write标志用于表示我们还没有初始化output，我们将使用这个张量去收集整个批量上的正确的检测。一旦进入了这个循环，我们需要清理一下。其中有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)

还记得我们已经将目标物体置信度小于阈值的边界框行置为0了吗？我们来清除一下吧。

        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.4if image_pred_.shape[0] == 0:continue

这个try-except代码块在这里是为了处理我们没有得到检测物体的情况。在这种情况下,我们使用continue去跳过这个图片的剩余循环体。

现在，让我们获得一种图片检测到的种类吧。

        #Get the various classes detected in the imageimg_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

然后，我们逐类使用非极大值抑制。img_classes中存放着检测到的类别。

        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
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

现在我们使用非极大值抑制

for i in range(idx):#Get the IOUs of all boxes that come after the one we are looking at #in the looptry:ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])except ValueError:breakexcept IndexError:break#Zero out all the detections that have IoU > treshholdiou_mask = (ious < nms_conf).float().unsqueeze(1)image_pred_class[i+1:] *= iou_mask       #Remove the non-zero entriesnon_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 non-zero entries
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind]

注意，我们已经把计算iou的代码行放在一个try-catch的代码块中了。这是因为循环用于运行idx迭代（image_pred_class中的行数）。然而，当我们继续循环的时候，一些边界框可能会从image_pred_class中移除出去。这就意味着，如果一个值从image_Pred_class中移除出去，我们就获取不到idx迭代。因此，我们可能尝试去索引一个值，这个值可能超出边界,或者image_pred_class[i+1:]可能返回一个空的张量，给它分配触发ValueErroe的错误。这个时候，我们可以确定NMS不能在移除边界框了，然后我们可以跳出这个循环。

3 计算IOU

这里是计算IOU的函数

def bbox_iou(box1, box2):"""Returns the IoU of two bounding boxes """#Get the coordinates of bounding boxesb1_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 rectangleinter_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 areainter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)#Union Areab1_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

4 编写预测

write_results函数输出一个尺寸为D*8的张量。这里D是所有图片中正确检测的数量，每一行代表一个。每一个检测有8个属性，也就是检测属于批量中的图片的索引，四个边角坐标，目标物体得分，最大置信度的种类的得分以及这些类的索引。像之前那样，除非我们已经有一个检测对象分配给输出向量了，否则我们不能够不初始化我们的输出张量。一旦它已经初始化了，我们就可以把后面的检测连接给它了。我们使用write标志去表示整个张量是否被初始化。在循环的最后迭代这个类，我们向output张量中添加相关的检测：

            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 imageseq = batch_ind, image_pred_classif not write:output = torch.cat(seq,1)write = Trueelse:out = torch.cat(seq,1)output = torch.cat((output,out))

在函数的最后，我们检查output是否已经初始化。如果没有初始化意味着，这里还没有批量中的任意一张图片的检测。这种情况，我们返回0：

    try:return outputexcept:return 0

这就是这篇文章的所有内容。在这篇文章的最后，我们最后得到了一个张量形式的预测，这个张量列出了所有的预测作为它的行。现在我们剩下的最后一件事就是，就是去创建一个输入的通道去从磁盘中读取图片，计算预测，在图片中画出bbox，并且展示/画出这些图片。这些我们将会在下一节中介绍。

5 总结：

在这篇文章中，主要对输入的图片进行目标物体得分阈值化和非极大值抑制。前者是为了选择最优可能出现目标物体的边界框，后者为对同一类的多个边界框先进行排除。这样就能够保证一个物体有一个边界框。在PyTorch中，主要有三个循环，第一个循环，迭代一个批量中所有的图片；第二个循环，迭代图片中所有的类别；第三个循环，迭代某个种类中所有的边界框，然后进行非极大值抑制。

（0）在迭代之前，我们先进行阈值化，将目标物体得分较低的边界框置零。

（1）进入第一次迭代，迭代每一张图片。

每个边界框，我们只关心边界框类别得分的最大值，以及最大值在原位置的索引，因此我们对输入张量进行改变，去掉80个类别信息，只保留边界框得分的最大值，以及其索引。并且将to列为零的边界框剔除（阈值化）。

注意：该图片可能并不包含物体（索引non_zero不存在或者[0]为0），因此我们使用continue来跳出本次循环的剩余内容，直接进行下一次的循环。

（2）进入第二次迭代，迭代某个类别的所有边界框。

在进入循环之前，我们需要得到目前图片中出现了哪些类别，然后再逐类进行迭代。在循环中，我们首先获得某一类的所有边界框，净所得到的边界框按照max_conf的数值从大到小进行排列。

（3）进入第三次迭代，依次对每张图片进行非极大值抑制

（4）在迭代每张图片所得到的结果，要存放在output张量中。

这个张量一共有八个属性：1个批次的序号，4个坐标值，1个目标物体置信度，1个最大类别得分，1个最大类别得分索引

注意：

在每次获得一个张量中特定的行的时候（阈值化，或者特定类别的行），我们可以使用mask和原张量进行相乘，再通过torch.nonzero().squeeze()得到特定行的索引，然后通过索引获得特定的行。
我们可以使用torch.cat()方法将几个元组或者列表连接，输出一个大的列表。