File size: 1,228 Bytes
186701e |
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 |
from typing import List, Tuple, Union
import cv2
from numpy import ndarray
MAJOR, MINOR = map(int, cv2.__version__.split('.')[:2])
assert MAJOR == 4
def non_max_suppression(boxes: Union[List[ndarray], Tuple[ndarray]],
scores: Union[List[float], Tuple[float]],
labels: Union[List[int], Tuple[int]],
conf_thres: float = 0.25,
iou_thres: float = 0.65) -> Tuple[List, List, List]:
if MINOR >= 7:
indices = cv2.dnn.NMSBoxesBatched(boxes, scores, labels, conf_thres,
iou_thres)
elif MINOR == 6:
indices = cv2.dnn.NMSBoxes(boxes, scores, conf_thres, iou_thres)
else:
indices = cv2.dnn.NMSBoxes(boxes, scores, conf_thres,
iou_thres).flatten()
nmsd_boxes = []
nmsd_scores = []
nmsd_labels = []
for idx in indices:
box = boxes[idx]
# x0y0wh -> x0y0x1y1
box[2:] = box[:2] + box[2:]
score = scores[idx]
label = labels[idx]
nmsd_boxes.append(box)
nmsd_scores.append(score)
nmsd_labels.append(label)
return nmsd_boxes, nmsd_scores, nmsd_labels
|