NMS
# -*- coding: utf-8 -*-
import cv2
import numpy as np
import random
def py_cpu_nms(dets, thresh):
"""Pure Python NMS baseline."""
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
scores = dets[:, 4] # 获取坐标和分数
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:]]) # 从1开始,第0个就是本身
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) # 重叠率,也就是IOU
inds = np.where(ovr <= thresh)[0] # 重叠率小于阈值的矩形框取出来
order = order[inds + 1] # np.where返回的是从零开始的,order是从一开始的,所以要加1
"""
[0, 1, 2, 3, 4, ]
[0.2, 0.4, 0.06, 0.1, 0.24]
概率由高到低排序
order [1, 4, 0, 3, 2 ]
prob [0.4, 0.24, 0.2, 0.1, 0.06]
最高概率框与其余框的各种计算(IOU)
order [1, 4, 0, 3, 2 ]
prob [0.4, 0.24, 0.2, 0.1, 0.06]
IOU [1, 0.7, 0.2, 0.6, 0.1]
取出IOU<thread(0.5),得到留下来的框
order' = order[1:]
id [0, 1, 2, 3. 4 ]
order [1, 4, 0, 3, 2 ]
id' [0, 1, 2, 3 ]
order' [4, 0, 3, 2 ]
IOU [0.7, 0.2, 0.6, 0.1 ]
inds' = [ 1, 3 ] # 对应id'
ids = [ 2, 4 ] # 对应id
# 新一轮的开始
order [4, 2 ]
prob [0.24, 0.06]
"""
return keep
def plot_rect(dets,img,name):
font = cv2.FONT_HERSHEY_SIMPLEX
for point in dets:
x1,y1,x2,y2,s=point
x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), thickness=2)
cv2.rectangle(img, (x1+2, y1+2), (x2, y1+30),(255, 0, 0), thickness=-1)
cv2.putText(img, str(s), (x1+2, y1+2), font, 0.6, (0, 0, 0), 1)
cv2.imshow(name,img)
cv2.waitKey(0)
def main():
img = np.zeros((1000,1000,3), np.uint8)
#fill the image with white
img.fill(255)
points = []
for _ in range(10):
x1 = random.randint(0,500)
y1 = random.randint(0,500)
x2 = random.randint(500,1000)
y2 = random.randint(500,1000)
s = random.random()
points.append([x1,y1,x2,y2,s])
points = np.array(points)
plot_rect(points,img,'1')
keep = py_cpu_nms(points,0.5)
dets = points[keep]
plot_rect(dets,img,'2')
cv2.destroyAllWindows()
main()