手撕NMS代码(python)
import numpy as np
def nms(dets,score):
x1 = dets[:,0]
y1 = dets[:,1]
x2 = dets[:,2]
y2 = dets[:,3]
area = (x2-x1+1)*(y2-y1+1)
scores = dets[:,4]
idx = scores.argsort()[::-1]
keep = []
while idx.size > 0:
i = idx[0]
keep.append(i)
xx1 = np.maximum(x1[i],x1[idx[1:]])
xx2 = np.minimum(x2[i],x2[idx[1:]])
yy1 = np.maximum(y1[i],y1[idx[1:]])
yy2 = np.minimum(y2[i],y2[idx[1:]])
w = np.maximum(xx2-xx1,0)
h = np.maximum(yy2-yy1,0)
overlap = w*h
iou = overlap/(area[i]+area[idx[1:]]-overlap)
idx_i = np.where(iou<=score_thr)[0]
idx = idx[idx_i+1]
return keep
手撕IOU代码
import numpy as np
def ComputeIOU(boxA, boxB):
## 计算相交框的坐标
x1 = np.max([boxA[0], boxB[0]])
x2 = np.min([boxA[2], boxB[2]])
y1 = np.max([boxA[1], boxB[1]])
y2 = np.min([boxA[3], boxB[3]])
## 计算交区域,并区域,及IOU
interArea = np.max([x2-x1+1, 0])*np.max([y2-y1+1,0]) ##一定要和0比较大小,如果是负数就说明压根不相交
unionArea = (boxA[2]-boxA[0]+1)*(boxA[3]-boxA[1]+1) + (boxB[2]-boxB[0]+1)*(boxB[3]-boxB[1]+1)-interArea
iou = interArea/unionArea
return iou
boxA = [1,1,3,3]
boxB = [2,2,4,4]
IOU = ComputeIOU(boxA, boxB)