mirror of
https://github.com/openmv/openmv.git
synced 2025-09-26 23:09:13 +08:00
scripts/libraries: Add keypoint support to NMS.
NMS will now pass-through keypoints along with scaling/offsetting them to be drawn correctly on the image.
This commit is contained in:
parent
eccb90699d
commit
7124b5134d
@ -43,7 +43,7 @@ class NMS:
|
|||||||
raise ValueError("Invalid ROI dimensions!")
|
raise ValueError("Invalid ROI dimensions!")
|
||||||
self.boxes = []
|
self.boxes = []
|
||||||
|
|
||||||
def add_bounding_box(self, xmin, ymin, xmax, ymax, score, label_index):
|
def add_bounding_box(self, xmin, ymin, xmax, ymax, score, label_index, keypoints=None):
|
||||||
if score >= 0.0 and score <= 1.0:
|
if score >= 0.0 and score <= 1.0:
|
||||||
xmin = max(0.0, min(xmin, self.window_w))
|
xmin = max(0.0, min(xmin, self.window_w))
|
||||||
ymin = max(0.0, min(ymin, self.window_h))
|
ymin = max(0.0, min(ymin, self.window_h))
|
||||||
@ -52,7 +52,7 @@ class NMS:
|
|||||||
w = int(xmax - xmin)
|
w = int(xmax - xmin)
|
||||||
h = int(ymax - ymin)
|
h = int(ymax - ymin)
|
||||||
if w > 0 and h > 0:
|
if w > 0 and h > 0:
|
||||||
self.boxes.append([int(xmin), int(ymin), w, h, score, label_index])
|
self.boxes.append([int(xmin), int(ymin), w, h, score, label_index, keypoints])
|
||||||
|
|
||||||
def get_bounding_boxes(self, threshold=0.1, sigma=0.1):
|
def get_bounding_boxes(self, threshold=0.1, sigma=0.1):
|
||||||
sorted_boxes = sorted(self.boxes, key=lambda x: x[4], reverse=True)
|
sorted_boxes = sorted(self.boxes, key=lambda x: x[4], reverse=True)
|
||||||
@ -107,15 +107,22 @@ class NMS:
|
|||||||
output_boxes[i][1] = int((output_boxes[i][1] * scale) + y_offset)
|
output_boxes[i][1] = int((output_boxes[i][1] * scale) + y_offset)
|
||||||
output_boxes[i][2] = int(output_boxes[i][2] * scale)
|
output_boxes[i][2] = int(output_boxes[i][2] * scale)
|
||||||
output_boxes[i][3] = int(output_boxes[i][3] * scale)
|
output_boxes[i][3] = int(output_boxes[i][3] * scale)
|
||||||
|
keypoints = output_boxes[i][6]
|
||||||
|
if keypoints is not None:
|
||||||
|
keypoints *= scale
|
||||||
|
keypoints[:, 0] += x_offset
|
||||||
|
keypoints[:, 1] += y_offset
|
||||||
|
|
||||||
# Create a list per class with (rect, score) tuples.
|
# Create a list per class with (rect, score) tuples.
|
||||||
|
|
||||||
output_list = [[] for i in range(max_label_index + 1)]
|
output_list = [[] for i in range(max_label_index + 1)]
|
||||||
|
|
||||||
for i in range(len(output_boxes)):
|
for i in range(len(output_boxes)):
|
||||||
output_list[output_boxes[i][5]].append(
|
rect_score = [output_boxes[i][:4], output_boxes[i][4]]
|
||||||
(output_boxes[i][0:4], output_boxes[i][4])
|
keypoints = output_boxes[i][6]
|
||||||
)
|
if keypoints is not None:
|
||||||
|
rect_score.append(keypoints)
|
||||||
|
output_list[output_boxes[i][5]].append(tuple(rect_score))
|
||||||
|
|
||||||
return output_list
|
return output_list
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user