IOU交并比和NMS非极大值抑制——pytorch代码实现

一、交并比 (IOU)

我们可以将其理解为在目标识别中的损失函数，损失函数的作用就是判断两个目标的距离或者相似程度。那么IOU就是判断两个方框之间的相似度。使用的是两个方框的交集除以并集。

这里需要注意的是：图片中的坐标和数学中不一样。x轴一样，y轴是相反的。

代码实现：

def IOU(box, other_boxes):
    #box: [x1,y1,x2,y2] 分别表示方框的左上角的点和右下角的点
    #other_boxs: N个box，多了一个维度(代表box的数量)
    box_area = (box[2]-box[0])*(box[3]-box[1])
    other_boxes_area = (other_boxes[:,2]-other_boxes[:,0]) * (other_boxes[:,3]-other_boxes[:,1])

    #交集
    x1 = torch.max(box[0],other_boxes[:,0])
    y1 = torch.max(box[1],other_boxes[:,1])
    x2 = torch.min(box[2],other_boxes[:,2])
    y2 = torch.min(box[3],other_boxes[:,3])
    Min = torch.tensor([0])
    w,h = torch.max(Min,x2-x1),torch.max(Min,y2-y1) #如果没有相交的框，两者相减是负值，为防止此类事情出现
    # 交集的面积
    overlap_area = w*h

    # 交并比 = 交集 / 并集
    iou = overlap_area / (box_area+other_boxes_area-overlap_area)
    return iou

二、非极大值抑制 (NMS)

由于使用图像金字塔和滑动窗口，我们的目标有可能出现下面的情况，如下图所示，一个人脸被多个框框选住。

这是我们可以使用非极大值抑制来实现对框的筛选。

• 第一步：找出所有预测目标中得分最高的预测框作为基准框；
• 第二步：计算剩余的预测框与基准框的IOU，如果IOU>阈值t，则将这些预测框删除，因为，这些预测框与0.97的太相似了，所以是冗余框，需要抑制；
• 第三步：在去除冗余框后剩余预测框中，找到得分第二高的预测框0.95作为基准框，执行第二步操作；
• 第四步：重复上述过程，直至剩余预测框中的每个预测框都曾被用作过基准框。即剩余的预测框任意两个IOU均小于t，没有两个预测框过于相似，完成nms操作。

def NMS(boxes, thre = 0.5):
    # boxes: [[置信度, x1,y1,x2,y2], [置信度, x1,y1,x2,y2], ...]
    # - 第一步：找出所有预测目标中得分最高的预测框作为基准框；
    # - 第二步：计算剩余的预测框与基准框的IOU，如果IOU>阈值t，则将这些预测框删除，因为，这些预测框与0.97的太相似了，所以是冗余框，需要抑制；
    # - 第三步：在去除冗余框后剩余预测框中，找到得分第二高的预测框0.95作为基准框，执行第二步操作；
    # - 第四步：重复上述过程，直至剩余预测框中的每个预测框都曾被用作过基准框。即剩余的预测框任意两个IOU均小于t，没有两个预测框过于相似，完成nms操作。
    sort_boxes = boxes[boxes[:, 0].argsort(descending=True)] # 获得置信度降序排序的下标，并根据下标进行索引。
    res = []

    # 防止传入错误的值
    if len(sort_boxes) == 0:
        print("ERROR: NO boxes! ")
        return []

    while len(sort_boxes) > 1:
        ref_box = sort_boxes[0]
        res.append(ref_box)

        other_boxes = sort_boxes[1:]
        sort_boxes = other_boxes[torch.where(IOU(ref_box[1:], other_boxes[:,1:])<thre)]
    
    if len(sort_boxes)>0: res.append(sort_boxes[0])

    return torch.stack(res)