【目标检测】 IoU 计算及 NMS 计算

IoU 计算

计算 IoU 的代码因为不是很复杂，所以在一些偏深度学习的岗位面试时比较容易遇到。一般都是给定一个点的形式为，[x1,y1,x2,y2] 或者是 [x1,y1,w1,h1]。

'''
Input: 
p_x=[x1,y1,w1,h1]
p_y=[x2,y2,w2,h2]

'''

def IoU(p_x,p_y):
    area_x = p_x[2]*p_x[3]
    area_y = p_y[2]*p_y[3]
    
    x_tl = max(p_x[0],p_y[0])
    y_tl = max(p_x[1],p_y[1])
    x_br = min(p_x[0]+p_x[2],p_y[0]+p_y[2])
    y_br = min(p_x[1]+p_x[3],p_y[1]+p_y[3])

    inter_w = max(x_br-x_tl,0)   # 与 0 比较是为了防止两个不相交的情况
    inter_h = max(y_br-y_tl,0)   # 与 0 比较是为了防止两个不相交的情况

    inter_area = inter_h*inter_w

    return inter_area*1.0 / (area_x+area_y - inter_area)

NMS 计算

虽然 NMS 需要基于上面的计算 IoU 公式，但是上面是简单情况。只有两个框，这里需要借助 Numpy，直接多维运算，否则程序写起来会很麻烦，运行起来也会因为多个 for 循环慢很多，这里直接参考：

import numpy as np

'''
input:
dets=[[x1,y1,x2,y2],[x1,y1,x2,y2],...] 
thresh: float
'''

def nms(dets, thresh):
   # 以下 5 个都是多维的数组
    x1 = dets[:, 0] 
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]
   
   # 用数组存储所有 box 的面积
    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    
    # 将数组的类别分数从大到小排序，并得到它们排序后的索引顺序 
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]  # np.where() 直接将小于阈值的去除了，返回 tuple（[index1,index2,...]）
        order = order[inds + 1]

    return keep