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