-
Notifications
You must be signed in to change notification settings - Fork 88
/
nms.py
119 lines (110 loc) · 4.08 KB
/
nms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import numpy as np
import cv2
def calc_IoU(a, b):
# step1:
inter_x1 = np.maximum(np.min(a[:,0]), np.min(b[:,0]))
inter_x2 = np.minimum(np.max(a[:,0]), np.max(b[:,0]))
inter_y1 = np.maximum(np.min(a[:,1]), np.min(b[:,1]))
inter_y2 = np.minimum(np.max(a[:,1]), np.max(b[:,1]))
if inter_x1>=inter_x2 or inter_y1>=inter_y2:
return 0.
x1 = np.minimum(np.min(a[:,0]), np.min(b[:,0]))
x2 = np.maximum(np.max(a[:,0]), np.max(b[:,0]))
y1 = np.minimum(np.min(a[:,1]), np.min(b[:,1]))
y2 = np.maximum(np.max(a[:,1]), np.max(b[:,1]))
if x1>=x2 or y1>=y2 or (x2-x1)<2 or (y2-y1)<2:
return 0.
else:
mask_w = np.int(np.ceil(x2-x1))
mask_h = np.int(np.ceil(y2-y1))
mask_a = np.zeros(shape=(mask_h, mask_w), dtype=np.uint8)
mask_b = np.zeros(shape=(mask_h, mask_w), dtype=np.uint8)
a[:,0] -= x1
a[:,1] -= y1
b[:,0] -= x1
b[:,1] -= y1
mask_a = cv2.fillPoly(mask_a, pts=np.asarray([a], 'int32'), color=1)
mask_b = cv2.fillPoly(mask_b, pts=np.asarray([b], 'int32'), color=1)
inter = np.logical_and(mask_a, mask_b).sum()
union = np.logical_or(mask_a, mask_b).sum()
iou = float(inter)/(float(union)+1e-12)
# print(iou)
# cv2.imshow('img1', np.uint8(mask_a*255))
# cv2.imshow('img2', np.uint8(mask_b*255))
# k = cv2.waitKey(0)
# if k==ord('q'):
# cv2.destroyAllWindows()
# exit()
return iou
def draw_image(pts, image):
cen_pts = np.mean(pts, axis=0)
tt = pts[0, :]
rr = pts[1, :]
bb = pts[2, :]
ll = pts[3, :]
cv2.line(image, (int(cen_pts[0]), int(cen_pts[1])), (int(tt[0]), int(tt[1])), (0, 0, 255), 2, 1)
cv2.line(image, (int(cen_pts[0]), int(cen_pts[1])), (int(rr[0]), int(rr[1])), (255, 0, 255), 2, 1)
cv2.line(image, (int(cen_pts[0]), int(cen_pts[1])), (int(bb[0]), int(bb[1])), (0, 255, 0), 2, 1)
cv2.line(image, (int(cen_pts[0]), int(cen_pts[1])), (int(ll[0]), int(ll[1])), (255, 0, 0), 2, 1)
return image
def NMS_numpy_exboxes(exboxes, conf, nms_thresh=0.5, image=None):
if len(exboxes)==0:
return None
sorted_index = np.argsort(conf) # Ascending order
keep_index = []
while len(sorted_index)>0:
curr_index = sorted_index[-1]
keep_index.append(curr_index)
if len(sorted_index)==1:
break
sorted_index = sorted_index[:-1]
IoU = []
for index in sorted_index:
iou = calc_IoU(exboxes[index,:,:].copy(), exboxes[curr_index,:,:].copy())
IoU.append(iou)
IoU = np.asarray(IoU, np.float32)
sorted_index = sorted_index[IoU<=nms_thresh]
return keep_index
def NMS_numpy_bbox(bboxes, nms_thresh=0.5):
"""
bboxes: num_insts x 5 [x1,y1,x2,y2,conf]
"""
if len(bboxes)==0:
return None
x1 = bboxes[:,0]
y1 = bboxes[:,1]
x2 = bboxes[:,2]
y2 = bboxes[:,3]
conf = bboxes[:,4]
area_all = (x2-x1)*(y2-y1)
sorted_index = np.argsort(conf) # Ascending order
keep_index = []
while len(sorted_index)>0:
# get the last biggest values
curr_index = sorted_index[-1]
keep_index.append(curr_index)
if len(sorted_index)==1:
break
# pop the value
sorted_index = sorted_index[:-1]
# get the remaining boxes
yy1 = np.take(y1, indices=sorted_index)
xx1 = np.take(x1, indices=sorted_index)
yy2 = np.take(y2, indices=sorted_index)
xx2 = np.take(x2, indices=sorted_index)
# get the intersection box
yy1 = np.maximum(yy1, y1[curr_index])
xx1 = np.maximum(xx1, x1[curr_index])
yy2 = np.minimum(yy2, y2[curr_index])
xx2 = np.minimum(xx2, x2[curr_index])
# calculate IoU
w = xx2-xx1
h = yy2-yy1
w = np.maximum(0., w)
h = np.maximum(0., h)
inter = w*h
rem_areas = np.take(area_all, indices=sorted_index)
union = (rem_areas-inter)+area_all[curr_index]
IoU = inter/union
sorted_index = sorted_index[IoU<=nms_thresh]
return keep_index