mirror of
https://github.com/openmv/openmv.git
synced 2025-09-26 23:09:13 +08:00
scripts/libraries: Expose nms settings in yolo post-processing callbacks.
This commit is contained in:
parent
7957589367
commit
6d21326233
@ -72,7 +72,7 @@ class yolo_v2_postprocess:
|
||||
_YOLO_V2_SCORE = const(4)
|
||||
_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
|
||||
if anchors is not None:
|
||||
self.anchors = anchors
|
||||
@ -83,6 +83,8 @@ class yolo_v2_postprocess:
|
||||
[5.55170, 9.30660],
|
||||
[9.72600, 11.1422]], dtype=np.float)
|
||||
self.anchors_len = len(self.anchors)
|
||||
self.nms_threshold = nms_threshold
|
||||
self.nms_sigma = nms_sigma
|
||||
|
||||
def __call__(self, model, inputs, outputs):
|
||||
ob, oh, ow, oc = model.output_shape[0]
|
||||
@ -152,7 +154,7 @@ class yolo_v2_postprocess:
|
||||
x_center[i] + (w_rel[i] / 2),
|
||||
y_center[i] + (h_rel[i] / 2),
|
||||
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:
|
||||
@ -163,8 +165,10 @@ class yolo_v5_postprocess:
|
||||
_YOLO_V5_SCORE = const(4)
|
||||
_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.nms_threshold = nms_threshold
|
||||
self.nms_sigma = nms_sigma
|
||||
|
||||
def __call__(self, model, inputs, outputs):
|
||||
oh, ow, oc = model.output_shape[0]
|
||||
@ -208,4 +212,4 @@ class yolo_v5_postprocess:
|
||||
for i in range(len(bb)):
|
||||
nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[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)
|
||||
|
Loading…
Reference in New Issue
Block a user