def test_iou():
box = np.array([
10, 20, 80, 60,
])
boxes = np.array([
[10, 20, 80, 60, ],
[10, 20, 40, 30, ],
[50, 50, 80, 60, ],
])
expected = np.array([
1,
0.25,
(40 * 30) / (80 * 60 * 2 - 40 * 30),
])
ious = iou(boxes, box)
assert np.allclose(ious, expected)
def _nms(self, boxes):
scores = boxes[:, 5]
order_indices = np.argsort(-scores)
keep_indices = []
while order_indices.size > 0:
i = order_indices[0]
keep_indices.append(i)
if order_indices.size == 1:
break
ious = iou(boxes[order_indices[1:], :], boxes[i, :])
remain_boxes = np.where(ious < self.iou_threshold)[0]
remain_boxes = remain_boxes + 1
order_indices = order_indices[remain_boxes]
nms_boxes = boxes[keep_indices, :]
if len(nms_boxes) > self.max_output_size:
nms_boxes = nms_boxes[:self.max_output_size, :]
return nms_boxes
