本篇文章是《如何使用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 = False
for 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.4
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
因为这里可能会出现一个类有多个正确检测到的物体,我们使用一个叫做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 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)
这里我们使用一个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 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 = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(inter_rect_y2 - inter_rect_y1 + 1, min=0)
#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
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 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))
在函数的最后,我们检查output是否已经初始化。如果没有初始化意味着,这里还没有批量中的任意一张图片的检测。这种情况,我们返回0:
try:
return output
except:
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()方法将几个元组或者列表连接,输出一个大的列表。