-
Notifications
You must be signed in to change notification settings - Fork 10
/
utils.py
128 lines (109 loc) · 5.41 KB
/
utils.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
120
121
122
123
124
125
126
127
128
import numpy as np
def decode_bbox(anchors, raw_outputs, variances=[0.1, 0.1, 0.2, 0.2]):
'''
Decode the actual bbox according to the anchors.
the anchor value order is:[xmin,ymin, xmax, ymax]
:param anchors: numpy array with shape [batch, num_anchors, 4]
:param raw_outputs: numpy array with the same shape with anchors
:param variances: list of float, default=[0.1, 0.1, 0.2, 0.2]
:return:
'''
anchor_centers_x = (anchors[:, :, 0:1] + anchors[:, :, 2:3]) / 2
anchor_centers_y = (anchors[:, :, 1:2] + anchors[:, :, 3:]) / 2
anchors_w = anchors[:, :, 2:3] - anchors[:, :, 0:1]
anchors_h = anchors[:, :, 3:] - anchors[:, :, 1:2]
raw_outputs_rescale = raw_outputs * np.array(variances)
predict_center_x = raw_outputs_rescale[:, :, 0:1] * anchors_w + anchor_centers_x
predict_center_y = raw_outputs_rescale[:, :, 1:2] * anchors_h + anchor_centers_y
predict_w = np.exp(raw_outputs_rescale[:, :, 2:3]) * anchors_w
predict_h = np.exp(raw_outputs_rescale[:, :, 3:]) * anchors_h
predict_xmin = predict_center_x - predict_w / 2
predict_ymin = predict_center_y - predict_h / 2
predict_xmax = predict_center_x + predict_w / 2
predict_ymax = predict_center_y + predict_h / 2
predict_bbox = np.concatenate([predict_xmin, predict_ymin, predict_xmax, predict_ymax], axis=-1)
return predict_bbox
def generate_anchors(feature_map_sizes, anchor_sizes, anchor_ratios, offset=0.5):
'''
generate anchors.
:param feature_map_sizes: list of list, for example: [[40,40], [20,20]]
:param anchor_sizes: list of list, for example: [[0.05, 0.075], [0.1, 0.15]]
:param anchor_ratios: list of list, for example: [[1, 0.5], [1, 0.5]]
:param offset: default to 0.5
:return:
'''
anchor_bboxes = []
for idx, feature_size in enumerate(feature_map_sizes):
cx = (np.linspace(0, feature_size[0] - 1, feature_size[0]) + 0.5) / feature_size[0]
cy = (np.linspace(0, feature_size[1] - 1, feature_size[1]) + 0.5) / feature_size[1]
cx_grid, cy_grid = np.meshgrid(cx, cy)
cx_grid_expend = np.expand_dims(cx_grid, axis=-1)
cy_grid_expend = np.expand_dims(cy_grid, axis=-1)
center = np.concatenate((cx_grid_expend, cy_grid_expend), axis=-1)
num_anchors = len(anchor_sizes[idx]) + len(anchor_ratios[idx]) - 1
center_tiled = np.tile(center, (1, 1, 2 * num_anchors))
anchor_width_heights = []
# different scales with the first aspect ratio
for scale in anchor_sizes[idx]:
ratio = anchor_ratios[idx][0] # select the first ratio
width = scale * np.sqrt(ratio)
height = scale / np.sqrt(ratio)
anchor_width_heights.extend([-width / 2.0, -height / 2.0, width / 2.0, height / 2.0])
# the first scale, with different aspect ratios (except the first one)
for ratio in anchor_ratios[idx][1:]:
s1 = anchor_sizes[idx][0] # select the first scale
width = s1 * np.sqrt(ratio)
height = s1 / np.sqrt(ratio)
anchor_width_heights.extend([-width / 2.0, -height / 2.0, width / 2.0, height / 2.0])
bbox_coords = center_tiled + np.array(anchor_width_heights)
bbox_coords_reshape = bbox_coords.reshape((-1, 4))
anchor_bboxes.append(bbox_coords_reshape)
anchor_bboxes = np.concatenate(anchor_bboxes, axis=0)
return anchor_bboxes
def single_class_non_max_suppression(bboxes, confidences, conf_thresh=0.2, iou_thresh=0.5, keep_top_k=-1):
'''
do nms on single class.
Hint: for the specific class, given the bbox and its confidence,
1) sort the bbox according to the confidence from top to down, we call this a set
2) select the bbox with the highest confidence, remove it from set, and do IOU calculate with the rest bbox
3) remove the bbox whose IOU is higher than the iou_thresh from the set,
4) loop step 2 and 3, util the set is empty.
:param bboxes: numpy array of 2D, [num_bboxes, 4]
:param confidences: numpy array of 1D. [num_bboxes]
:param conf_thresh:
:param iou_thresh:
:param keep_top_k:
:return:
'''
if len(bboxes) == 0: return []
conf_keep_idx = np.where(confidences > conf_thresh)[0]
bboxes = bboxes[conf_keep_idx]
confidences = confidences[conf_keep_idx]
pick = []
xmin = bboxes[:, 0]
ymin = bboxes[:, 1]
xmax = bboxes[:, 2]
ymax = bboxes[:, 3]
area = (xmax - xmin + 1e-3) * (ymax - ymin + 1e-3)
idxs = np.argsort(confidences)
while len(idxs) > 0:
last = len(idxs) - 1
i = idxs[last]
pick.append(i)
# keep top k
if keep_top_k != -1:
if len(pick) >= keep_top_k:
break
overlap_xmin = np.maximum(xmin[i], xmin[idxs[:last]])
overlap_ymin = np.maximum(ymin[i], ymin[idxs[:last]])
overlap_xmax = np.minimum(xmax[i], xmax[idxs[:last]])
overlap_ymax = np.minimum(ymax[i], ymax[idxs[:last]])
overlap_w = np.maximum(0, overlap_xmax - overlap_xmin)
overlap_h = np.maximum(0, overlap_ymax - overlap_ymin)
overlap_area = overlap_w * overlap_h
overlap_ratio = overlap_area / (area[idxs[:last]] + area[i] - overlap_area)
need_to_be_deleted_idx = np.concatenate(([last], np.where(overlap_ratio > iou_thresh)[0]))
idxs = np.delete(idxs, need_to_be_deleted_idx)
# if the number of final bboxes is less than keep_top_k, we need to pad it.
# TODO
return conf_keep_idx[pick]