原文地址:https://blog.paperspace.com/how-to-implement-a-yolo-object-detector-in-pytorch/
代码:https://github.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch
在前面的部分中,我们构建了一个模型,该模型在给定输入图像的情况下输出多个对象检测。 确切地说,我们的输出是B x 10647 x 85的张量.B是批量图像的数量(batch内图像的数量),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
执行非最大抑制
我们现在拥有的边界框属性由中心坐标以及边界框的高度和宽度来描述。 但是,使用每个box的一对对角的坐标((coordinates of a pair of diagnal corners of each box)来计算两个box的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的批次,其中图像1,2和3分别具有5,2,4个真实检测。 因此,必须一次对一个图像进行置信度阈值处理和NMS。 这意味着,我们无法对所涉及的操作进行矢量化???,并且必须遍历预测的第一维(包含批处理中的图像索引)。
batch_size = prediction.size(0)write = Falsefor ind in range(batch_size):image_pred = prediction[ind] #image Tensor#confidence threshholding #NMS
如前所述,写标志用于指示我们尚未初始化输出,我们将使用张量来收集整个批次的真实检测(true detections)。
一旦进入循环,让我们清理一下。 请注意,每个边界框行有85(x,y,w,h,c,classes)个属性,其中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)
还记得我们已经将具有小于阈值的对象置信度的边界框行设置为零吗? 让我们去掉他们。
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块???用于处理我们没有检测到的情况。 在这种情况下,我们使用继续跳过此图像的循环体的其余部分。
现在,让我们在图像中检测出类。
#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
然后,我们按类别执行NMS。
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
执行NMS。
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的同一类的边界框,则消除具有较低类置信度的边界框。 我们已经整理出了具有更高置信度的边界框。
在循环体中,下面的行给出了box的IoU,由i索引,所有边界框的索引都高于i。
ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
每次迭代,如果具有大于i的索引的任何边界框具有大于阈值nms_thresh的IoU(具有由i索引的框),则消除该特定框。
#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]
还要注意,我们已经在try-catch块中放置了代码行来计算ious。 这是因为循环被设计为运行idx迭代(image_pred_class中的行数)。 但是,当我们继续循环时,可以从image_pred_class中删除许多边界框。 这意味着,即使从image_pred_class中删除了一个值,我们也无法进行idx迭代。 因此,我们可能会尝试索引超出边界的值(IndexError),或者切片image_pred_class [i + 1:]可能返回一个空张量,分配触发ValueError。 此时,我们可以确定NMS不能删除任何进一步的边界框,并且我们打破了循环。(Also notice, we have put the line of code to compute the ious
in a try-catch block. This is because the loop is designed to run idx
iterations (number of rows in image_pred_class
). However, as we proceed with the loop, a number of bounding boxes may be removed from image_pred_class
. This means, even if one value is removed from image_pred_class
, we cannot have idx
iterations. Hence, we might try to index a value that is out of bounds (IndexError
), or the slice image_pred_class[i+1:]
may return an empty tensor, assigning which triggers a ValueError
. At that point, we can ascertain that NMS can remove no further bounding boxes, and we break out of the loop.)
计算IoU
这是函数bbox_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
写预测
函数write_results输出形状D x 8的张量。这里D是所有图像中的真实检测(true detections),每个图像由一行表示。 每个检测具有8个属性,即检测所属批次中的图像的索引,4个角坐标,对象得分(objectness score),具有最大置信度的类的得分,以及该类的索引。
和以前一样,我们不会初始化输出张量,除非我们有一个检测分配给它。 一旦初始化,我们将后续检测连接到它。 我们使用写标志来指示张量是否已经初始化。 在遍历类的循环结束时,我们将结果检测添加到张量输出。
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))
在函数结束时,我们检查输出是否已完全初始化。 如果没有手段,则批次的任何图像中都没有单一检测。 在那种情况下,我们返回0。
try:return outputexcept:return 0
这是这篇文章的内容。 在这篇文章的最后,我们最终得到了一个张量形式的预测,它将每个预测列为行。 现在唯一剩下的就是创建一个输入管道来从磁盘读取图像,计算预测,在图像上绘制边界框,然后显示/写入这些图像。 这是我们将在下一部分中做的。