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_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)
|
||||||
|
Loading…
Reference in New Issue
Block a user