Create box_utils.py
Browse files- box_utils.py +115 -0
box_utils.py
ADDED
@@ -0,0 +1,115 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# SPDX-License-Identifier: MIT
|
2 |
+
|
3 |
+
import numpy as np
|
4 |
+
|
5 |
+
def area_of(left_top, right_bottom):
|
6 |
+
"""
|
7 |
+
Compute the areas of rectangles given two corners.
|
8 |
+
Args:
|
9 |
+
left_top (N, 2): left top corner.
|
10 |
+
right_bottom (N, 2): right bottom corner.
|
11 |
+
Returns:
|
12 |
+
area (N): return the area.
|
13 |
+
"""
|
14 |
+
hw = np.clip(right_bottom - left_top, 0.0, None)
|
15 |
+
return hw[..., 0] * hw[..., 1]
|
16 |
+
|
17 |
+
def iou_of(boxes0, boxes1, eps=1e-5):
|
18 |
+
"""
|
19 |
+
Return intersection-over-union (Jaccard index) of boxes.
|
20 |
+
Args:
|
21 |
+
boxes0 (N, 4): ground truth boxes.
|
22 |
+
boxes1 (N or 1, 4): predicted boxes.
|
23 |
+
eps: a small number to avoid 0 as denominator.
|
24 |
+
Returns:
|
25 |
+
iou (N): IoU values.
|
26 |
+
"""
|
27 |
+
overlap_left_top = np.maximum(boxes0[..., :2], boxes1[..., :2])
|
28 |
+
overlap_right_bottom = np.minimum(boxes0[..., 2:], boxes1[..., 2:])
|
29 |
+
|
30 |
+
overlap_area = area_of(overlap_left_top, overlap_right_bottom)
|
31 |
+
area0 = area_of(boxes0[..., :2], boxes0[..., 2:])
|
32 |
+
area1 = area_of(boxes1[..., :2], boxes1[..., 2:])
|
33 |
+
return overlap_area / (area0 + area1 - overlap_area + eps)
|
34 |
+
|
35 |
+
def hard_nms(box_scores, iou_threshold, top_k=-1, candidate_size=200):
|
36 |
+
"""
|
37 |
+
Perform hard non-maximum-supression to filter out boxes with iou greater
|
38 |
+
than threshold
|
39 |
+
Args:
|
40 |
+
box_scores (N, 5): boxes in corner-form and probabilities.
|
41 |
+
iou_threshold: intersection over union threshold.
|
42 |
+
top_k: keep top_k results. If k <= 0, keep all the results.
|
43 |
+
candidate_size: only consider the candidates with the highest scores.
|
44 |
+
Returns:
|
45 |
+
picked: a list of indexes of the kept boxes
|
46 |
+
"""
|
47 |
+
scores = box_scores[:, -1]
|
48 |
+
boxes = box_scores[:, :-1]
|
49 |
+
picked = []
|
50 |
+
indexes = np.argsort(scores)
|
51 |
+
indexes = indexes[-candidate_size:]
|
52 |
+
while len(indexes) > 0:
|
53 |
+
current = indexes[-1]
|
54 |
+
picked.append(current)
|
55 |
+
if 0 < top_k == len(picked) or len(indexes) == 1:
|
56 |
+
break
|
57 |
+
current_box = boxes[current, :]
|
58 |
+
indexes = indexes[:-1]
|
59 |
+
rest_boxes = boxes[indexes, :]
|
60 |
+
iou = iou_of(
|
61 |
+
rest_boxes,
|
62 |
+
np.expand_dims(current_box, axis=0),
|
63 |
+
)
|
64 |
+
indexes = indexes[iou <= iou_threshold]
|
65 |
+
|
66 |
+
return box_scores[picked, :]
|
67 |
+
|
68 |
+
def predict(width, height, confidences, boxes, prob_threshold, iou_threshold=0.5, top_k=-1):
|
69 |
+
"""
|
70 |
+
Select boxes that contain human faces
|
71 |
+
Args:
|
72 |
+
width: original image width
|
73 |
+
height: original image height
|
74 |
+
confidences (N, 2): confidence array
|
75 |
+
boxes (N, 4): boxes array in corner-form
|
76 |
+
iou_threshold: intersection over union threshold.
|
77 |
+
top_k: keep top_k results. If k <= 0, keep all the results.
|
78 |
+
Returns:
|
79 |
+
boxes (k, 4): an array of boxes kept
|
80 |
+
labels (k): an array of labels for each boxes kept
|
81 |
+
probs (k): an array of probabilities for each boxes being in corresponding labels
|
82 |
+
"""
|
83 |
+
boxes = boxes[0]
|
84 |
+
confidences = confidences[0]
|
85 |
+
#print(boxes)
|
86 |
+
#print(confidences)
|
87 |
+
|
88 |
+
picked_box_probs = []
|
89 |
+
picked_labels = []
|
90 |
+
for class_index in range(1, confidences.shape[1]):
|
91 |
+
#print(confidences.shape[1])
|
92 |
+
probs = confidences[:, class_index]
|
93 |
+
#print(probs)
|
94 |
+
mask = probs > prob_threshold
|
95 |
+
probs = probs[mask]
|
96 |
+
|
97 |
+
if probs.shape[0] == 0:
|
98 |
+
continue
|
99 |
+
subset_boxes = boxes[mask, :]
|
100 |
+
#print(subset_boxes)
|
101 |
+
box_probs = np.concatenate([subset_boxes, probs.reshape(-1, 1)], axis=1)
|
102 |
+
box_probs = hard_nms(box_probs,
|
103 |
+
iou_threshold=iou_threshold,
|
104 |
+
top_k=top_k,
|
105 |
+
)
|
106 |
+
picked_box_probs.append(box_probs)
|
107 |
+
picked_labels.extend([class_index] * box_probs.shape[0])
|
108 |
+
if not picked_box_probs:
|
109 |
+
return np.array([]), np.array([]), np.array([])
|
110 |
+
picked_box_probs = np.concatenate(picked_box_probs)
|
111 |
+
picked_box_probs[:, 0] *= width
|
112 |
+
picked_box_probs[:, 1] *= height
|
113 |
+
picked_box_probs[:, 2] *= width
|
114 |
+
picked_box_probs[:, 3] *= height
|
115 |
+
return picked_box_probs[:, :4].astype(np.int32), np.array(picked_labels), picked_box_probs[:, 4]
|