120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
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
|