这是关于从头实现YOLO v3检测器教程的第4部分。最后,我们实现了网络的转发。在这一部分中,我们使用对象置信度作为检测的阈值,然后使用非最大抑制。
本教程的代码设计为在Python 3.5和PyTorch 0.4上运行。教程完整代码可以在这里找到 Github repo.
本教程分为5个部分:
- Part 1 : Understanding How YOLO works
- Part 2 : Creating the layers of the network architecture
- Part 3 : Implementing the the forward pass of the network
- Part 4 (This one): Confidence Thresholding and Non-maximum Suppression
- Part 5 : Designing the input and the output pipelines
先修课程
- 课程1-3部分
- 基本的PyTorch工作知识,包括如何使用nn.Module、nn.sequence和torch.nn.parameter类创建自定义架构。
- NumPy基本知识
在前面的部分中,我们构建了一个模型,该模型在给定输入图像的情况下输出多个目标检测。准确地说,我们的输出是一个B×10647×85的张量。B为批处理图像的数量,10647为每幅图像预测的包围框的数量,85为包围框属性的数量。
但是,如第1部分所述,我们必须使输出符合对象分数阈值和非最大抑制,以获得我将在本文的其余部分中称为真实检测的内容。 为此,我们将在文件util.py中创建一个名为write_results的函数
def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
这些函数将预测、置信度(对象得分阈值)、num_classes(在我们的例子中是80)和nms_conf (NMS IoU阈值)作为输入。
对象置信度阈值
我们的预测张量包含有关B x 10647边界框的信息。 对于每个对象得分低于阈值的边界框,我们将其每个属性(表示边界框的整行)的值设置为零。
conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
prediction = prediction*conf_mask
执行非最大抑制
Note: I assume you understand what IoU (Intersection over union) is, and what Non-maximum suppression is. If that is not the case, refer to links at the end of the post).
我们现在拥有的边界框属性是由中心坐标以及边界框的高度和宽度描述的。然而,使用每个盒子对角的坐标计算两个盒子的IoU更容易。因此,我们将box的(center x, center y, height, width)转换为(top-left corner x, top-left corner y, right-bottom corner x, right-bottom corner 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的批次,其中图像1,2和3分别具有5,2,4个真实检测。 因此,必须同时对一个图像进行置信度阈值处理和NMS。 这意味着,我们无法对所涉及的操作进行矢量化,并且必须遍历prediction的第一维(包含批处理中的图像索引)。
batch_size = prediction.size(0)
write = False
for ind in range(batch_size):
image_pred = prediction[ind] #图片张量
#置信度阈值
#NMS
正如前面所描述的,write标志用于指示我们还没有初始化output,我们将使用这个张量在整个批处理中收集true的检测结果。
一旦进入循环,让我们清理一下。注意,每个边界框行有85个属性,其中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
#PyTorch 0.4 兼容
#由于上述代码没有引发异常,所以没有检测
#PyTorch 0.4中支持标量
if image_pred_.shape[0] == 0:
continue
try-except块用于处理 我们没有检测到的情况。在这种情况下,我们使用continue来跳过此映像的循环主体的其余部分。
现在,让我们在图像中检测类。
#获取图像中检测到的各种类
img_classes = unique(image_pred_[:,-1]) #-1索引保存类索引
因为同一个类可以有多个true的检测,所以我们使用一个名为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表示)。
以下代码在原始代码文件中缩进了三个块,但我没有在此处缩进,因为此页面上的空间有限。
#得到一个特定类的检测
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)
#对探测结果进行排序,使得具有最大对象置信度的条目位于顶部
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) #检测数量
现在,我们执行NMS。
for i in range(idx):
#得到我们在循环中看到的边框与后面的所有边框的IoU
try:
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
except ValueError:
break
except IndexError:
break
#将所有IoU > treshhold的检测值归零
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#删除非零项
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,第二个输入包含每个边界框。
如果我们有两个具有大于阈值的IoU的同一类的边界框,则消除具有较低类置信度的边界框。 我们已经整理出了具有更高置信度的边界框。
在循环体中,下面的行给出了边界框的IoU,由i索引,所有边界框的索引都高于i。
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
每次迭代中,如果任何索引值大于i的边界框具有大于阈值nms_thresh的IoU(框由i索引),则消除该特定框。
#将所有IoU > treshhold的检测值归零
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i+1:] *= iou_mask
#删除非零项
non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind]
还要注意,我们已经在try-catch块中放置了代码行来计算ious。 这是因为循环被设计成运行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
"""
#获取边界框坐标
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]
#获取交叉矩形坐标
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)
#交叉区域
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
#联合区域
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
Writing the predictions
函数write_results输出一个形状为Dx8的张量。这里D是所有图像中的真实检测值,每一个都用一行表示。每个检测有8个属性,即检测所属批次中的图像的索引、4个角坐标、物体度得分、最大置信度类的得分、该类的索引。
和以前一样,我们不初始化输出张量,除非我们有一个检测要分配给它。一旦初始化,我们将后续的检测连接到它。我们使用写标记来指示张量是否已经初始化。在遍历类的循环结束时,我们将检测结果添加到张量output。
batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)
#对于图像中类cls尽可能多的检测,重复batch_id
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))
在函数结束时,我们检查output是否已经初始化。如果没有,那就意味着这批次的任何图像中都没有检测到。在这种情况下,返回0。
try:
return output
except:
return 0
这就是本文的内容。在这篇文章的最后,我们终于有了一个张量形式的预测,它列出了每一个预测作为它的行。现在剩下的惟一一件事就是创建一个输入管道,从磁盘读取图像、计算预测、在图像上绘制边框,然后显示/写入这些图像。这就是我们下一部分要做的 part.
扩展阅读
- PyTorch tutorial
- IoU
- Non maximum suppresion
- Non-maximum Suppression