前一节我们实现了网络的前向传播。这一节我们对检测输出设置目标置信度阈值和进行非极大值抑制。
必要条件&＃xff1a;
1.此系列教程的Part1到Part3。
2.Pytorch的基本知识&＃xff0c;包括如何使用nn.Module&＃xff0c;nn.Sequential&＃xff0c;torch.nn.parameter类构建常规的结构
3.numpy的基础知识
此前我们已经建立了一个模型&＃xff0c;给定一张输入图片它能产生B*10674*85维的输出向量。B是批中图片的数目&＃xff0c;10674是每张图片预测的边界框数目&＃xff0c;85是边界框属性数目。
但就像我们在part1中描述的那样&＃xff0c;我们必须对输出进行目标置信度阈值化和非极大值抑制&＃xff0c;以获得最终剩余的真正检测。为此&＃xff0c;我们将在文件util.py中创建一个名为write_results的函数。
def write_results(prediction, confidence, num_classes, nms_conf &＃61; 0.4):
这个函数将prediction、confidence(目标置信度阈值)、num_classes(在我们的示例中是80)和nms_conf (NMS IoU的阈值)作为输入。
目标置信度阈值&＃xff1a;
我们的预测张量包含了关于B x 10647个边界框的信息。对于每个目标置信度低于阈值的边界框&＃xff0c;我们将它的每个属性(行向量)的值设置为零。
conf_mask &＃61; (prediction[:,:,4] > confidence).float().unsqueeze(2)
prediction&＃61; prediction*conf_mask
执行极大值抑制&＃xff1a;
我们现在拥有的是边界框的中心坐标以及高度和宽度&＃xff0c;然而使用边界框的对角点更容易计算IOU。因此&＃xff0c;我们将框的(center x, center y, height, width)属性转换为(左上角x&＃xff0c;左上角y&＃xff0c;右下角x&＃xff0c;右下角y)。
box_corner &＃61;prediction.new(prediction.shape)
box_corner[:,:,0]&＃61; (prediction[:,:,0] - prediction[:,:,2]/2)
box_corner[:,:,1] &＃61; (prediction[:,:,1] - prediction[:,:,3]/2)
box_corner[:,:,2] &＃61; (prediction[:,:,0] &＃43; prediction[:,:,2]/2)
box_corner[:,:,3] &＃61; (prediction[:,:,1] &＃43; prediction[:,:,3]/2)
prediction[:,:,:4] &＃61; box_corner[:,:,:4]
每幅图像中真实检测框的数目可能不同。例如&＃xff0c;一批大小为3的图像&＃xff0c;其中图像1、2和3分别有5、2、4个真检测值。因此&＃xff0c;每次必须对同一个图像进行置信阈值和NMS&＃xff0c;而不能对所涉及的操作进行矢量化&＃xff0c;必须在预测的第一个维度(包含成批图像的索引)上进行遍历操作。
batch_size &＃61;prediction.size(0)
write&＃61;Falsefor ind inrange(batch_size):
image_pred&＃61; prediction[ind] #image Tensor
#confidence threshholding
#NMS
write标志位用于指示我们是否对output进行了初始化&＃xff0c;将会使用一个向量来收集整个批中真实的预测。
循环的开始我们进行数据清理。因为每个边界框行有85个属性&＃xff0c;其中80个是类得分。我们只关心类得分最大值的那个&＃xff0c;所以会从每行中删除80个类得分&＃xff0c;添加具有最大值的类的索引&＃xff0c;以及该类的类得分。
max_conf, max_conf_score &＃61; torch.max(image_pred[:,5:5&＃43; num_classes], 1)
max_conf&＃61; max_conf.float().unsqueeze(1)
max_conf_score&＃61; max_conf_score.float().unsqueeze(1)
seq&＃61; (image_pred[:,:5], max_conf, max_conf_score)
image_pred&＃61; torch.cat(seq, 1)
我们前面已经将目标置信度得分低于阈值的边界框行属性设置为了0&＃xff0c;现在就筛除它们。
non_zero_ind &＃61; (torch.nonzero(image_pred[:,4]))try:
image_pred_&＃61; 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.4
if image_pred_.shape[0] &＃61;&＃61;0:continue
try-except块用于处理没有检测到的情况。在这种情况下&＃xff0c;我们使用continue跳过此图像的其余循环体。
接下来&＃xff0c;让我们在图像中检测目标。
#Get the various classes detected in the image
img_classes &＃61; unique(image_pred_[:,-1]) #-1 index holds the class index
因为对于同一个类别可能存在多个正确检测&＃xff0c;我们使用一个叫unique的函数来获得给定图片中所有出现的类。
defunique(tensor):
tensor_np&＃61;tensor.cpu().numpy()
unique_np&＃61;np.unique(tensor_np)
unique_tensor&＃61;torch.from_numpy(unique_np)
tensor_res&＃61;tensor.new(unique_tensor.shape)
tensor_res.copy_(unique_tensor)return tensor_res
之后我们对于每个类进行NMS
for cls inimg_classes:#perform NMS
一进入这个循环&＃xff0c;首先要做的事情就是提取对于某一特定类别的检测(用变量cls表示)
#get the detections with one particular class
cls_mask &＃61; image_pred_*(image_pred_[:,-1] &＃61;&＃61; cls).float().unsqueeze(1)
class_mask_ind&＃61; torch.nonzero(cls_mask[:,-2]).squeeze()
image_pred_class&＃61; 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 &＃61; torch.sort(image_pred_class[:,4], descending &＃61; True )[1]
image_pred_class&＃61;image_pred_class[conf_sort_index]
idx&＃61; image_pred_class.size(0) #Number of detections
然后我们进行NMS
for i inrange(idx):#Get the IOUs of all boxes that come after the one we are looking at
#in the loop
try:
ious&＃61; bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i&＃43;1:])exceptValueError:break
exceptIndexError:break
#Zero out all the detections that have IoU > treshhold
iou_mask &＃61; (ious image_pred_class[i&＃43;1:] *&＃61;iou_mask#Remove the non-zero entries
non_zero_ind &＃61; torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class&＃61; image_pred_class[non_zero_ind].view(-1,7)
这里我们使用到了一个函数bbox_iou。第一个输入参数是循环体变量i索引处的边界框&＃xff0c;第二个输入参数是多行边界框的一个tensor。函数bbox_iou的输出是一个tensor它包含了第一个输入的边界框与第二个输入的所有边界框的IOU。如下&＃xff1a;
![]()
之前我们已经将目标置信度高的边界框放在前面&＃xff0c;如果后面的边界框IoU值与前面的相比超过了阈值&＃xff0c;那后者就会被删去。
循环体里面下面这行计算IoU。
ious &＃61; bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i&＃43;1:])
每轮迭代&＃xff0c;如果有任何索引大于i的边界框与第i个边界框的IoU大于阈值nms_thresh&＃xff0c;那这个边界框就会被删除。
#Zero out all the detections that have IoU > treshhold
iou_mask &＃61; (ious image_pred_class[i&＃43;1:] *&＃61;iou_mask#Remove the non-zero entries
non_zero_ind &＃61; torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class&＃61; image_pred_class[non_zero_ind]
还要注意的是&＃xff0c;我们将计算IoU的代码行放在try-catch块中。这是因为此循环按照id进行迭代(image_pred_class中的行数)。但因为我们循环过程中可能会从image_pred_class中删除一些边界框。这样一来&＃xff0c;迭代可能会出现索引越界触发IndexError或者image_pred_class[i&＃43;1:]返回一个空张量触发ValueError。此时我们可以确定NMS已经无法删除多余的边界框了&＃xff0c;从而跳出循环。
计算IoU:
defbbox_iou(box1, box2):"""Returns the IoU of two bounding boxes"""
#Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 &＃61; box1[:,0], box1[:,1], box1[:,2], box1[:,3]
b2_x1, b2_y1, b2_x2, b2_y2&＃61; box2[:,0], box2[:,1], box2[:,2], box2[:,3]#get the corrdinates of the intersection rectangle
inter_rect_x1 &＃61;torch.max(b1_x1, b2_x1)
inter_rect_y1&＃61;torch.max(b1_y1, b2_y1)
inter_rect_x2&＃61;torch.min(b1_x2, b2_x2)
inter_rect_y2&＃61;torch.min(b1_y2, b2_y2)#Intersection area
inter_area &＃61; torch.clamp(inter_rect_x2 - inter_rect_x1 &＃43; 1, min&＃61;0) * torch.clamp(inter_rect_y2 - inter_rect_y1 &＃43; 1, min&＃61;0)#Union Area
b1_area &＃61; (b1_x2 - b1_x1 &＃43; 1)*(b1_y2 - b1_y1 &＃43; 1)
b2_area&＃61; (b2_x2 - b2_x1 &＃43; 1)*(b2_y2 - b2_y1 &＃43; 1)
iou&＃61; inter_area / (b1_area &＃43; b2_area -inter_area)return iou
写入预测&＃xff1a;
write_results函数输出一个形状为 Dx8 的tensor。这里D是所有图像的真实检测&＃xff0c;每个都用一行表示。每个检测有8个属性&＃xff0c;即检测所属批次图像的索引、4个角坐标、目标置信度得分、最大置信类得分、该类的索引。
和此前一样&＃xff0c;我们等到有一个检测时才初始化输出向量并将后续的检测拼接进来。使用写标志来表示tensor是否已经初始化。在遍历类的循环结束时&＃xff0c;我们将检测结果添加到输出tensor中。
batch_ind &＃61; 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 &＃61;batch_ind, image_pred_classif notwrite:
output&＃61; torch.cat(seq,1)
write&＃61;Trueelse:
out&＃61; torch.cat(seq,1)
output&＃61; torch.cat((output,out))
在函数的末尾&＃xff0c;我们检查输出是否已经初始化。如果没有&＃xff0c;就意味着这批图像中没有一个检测到。在这种情况下&＃xff0c;我们返回0。
try:returnoutputexcept:
return 0
这就是这一部分所要讲解的内容了。现在我们终于有了一个预测&＃xff0c;它以tensor的形式列出了每一个边界框。所以只剩下一件事就是创建一个输入管道来从磁盘读取图像&＃xff0c;计算预测&＃xff0c;在图像上绘制边界框&＃xff0c;然后显示/写入这些图像。这是我们下一部分要做的。
Further Reading
京公网安备 11010802041100号