scripts/libraries: Expose nms settings in yolo post-processing callbacks.

This commit is contained in:
Kwabena W. Agyeman 2025-01-21 22:08:45 -08:00
parent 7957589367
commit 6d21326233

View File

@ -72,7 +72,7 @@ class yolo_v2_postprocess:
_YOLO_V2_SCORE = const(4) _YOLO_V2_SCORE = const(4)
_YOLO_V2_CLASSES = const(5) _YOLO_V2_CLASSES = const(5)
def __init__(self, threshold=0.6, anchors=None): def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold self.threshold = threshold
if anchors is not None: if anchors is not None:
self.anchors = anchors self.anchors = anchors
@ -83,6 +83,8 @@ class yolo_v2_postprocess:
[5.55170, 9.30660], [5.55170, 9.30660],
[9.72600, 11.1422]], dtype=np.float) [9.72600, 11.1422]], dtype=np.float)
self.anchors_len = len(self.anchors) self.anchors_len = len(self.anchors)
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
ob, oh, ow, oc = model.output_shape[0] ob, oh, ow, oc = model.output_shape[0]
@ -152,7 +154,7 @@ class yolo_v2_postprocess:
x_center[i] + (w_rel[i] / 2), x_center[i] + (w_rel[i] / 2),
y_center[i] + (h_rel[i] / 2), y_center[i] + (h_rel[i] / 2),
bb_scores[i], bb_classes[i]) bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes() return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)
class yolo_v5_postprocess: class yolo_v5_postprocess:
@ -163,8 +165,10 @@ class yolo_v5_postprocess:
_YOLO_V5_SCORE = const(4) _YOLO_V5_SCORE = const(4)
_YOLO_V5_CLASSES = const(5) _YOLO_V5_CLASSES = const(5)
def __init__(self, threshold=0.6): def __init__(self, threshold=0.6, nms_threshold=0.1, nms_sigma=0.1):
self.threshold = threshold self.threshold = threshold
self.nms_threshold = nms_threshold
self.nms_sigma = nms_sigma
def __call__(self, model, inputs, outputs): def __call__(self, model, inputs, outputs):
oh, ow, oc = model.output_shape[0] oh, ow, oc = model.output_shape[0]
@ -208,4 +212,4 @@ class yolo_v5_postprocess:
for i in range(len(bb)): for i in range(len(bb)):
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i],
bb_scores[i], bb_classes[i]) bb_scores[i], bb_classes[i])
return nms.get_bounding_boxes() return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)