mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
scripts/libraries: Cleanup yolo v2 post-processing.
This commit is contained in:
parent
2b98a4e963
commit
d67bd1ad7c
@ -33,6 +33,9 @@ from micropython import const
|
|||||||
from ulab import numpy as np
|
from ulab import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
_NO_DETECTION = const(())
|
||||||
|
|
||||||
|
|
||||||
# FOMO generates an image per class, where each pixel represents the centroid
|
# FOMO generates an image per class, where each pixel represents the centroid
|
||||||
# of the trained object. These images are processed with `find_blobs()` to
|
# of the trained object. These images are processed with `find_blobs()` to
|
||||||
# extract centroids, and `get_stats()` is used to get their scores. Overlapping
|
# extract centroids, and `get_stats()` is used to get their scores. Overlapping
|
||||||
@ -69,8 +72,8 @@ 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, score_threshold=0.6, anchors=None):
|
def __init__(self, threshold=0.6, anchors=None):
|
||||||
self.score_threshold = score_threshold
|
self.threshold = threshold
|
||||||
if anchors is not None:
|
if anchors is not None:
|
||||||
self.anchors = anchors
|
self.anchors = anchors
|
||||||
else:
|
else:
|
||||||
@ -101,11 +104,11 @@ class yolo_v2_postprocess:
|
|||||||
|
|
||||||
# Threshold all the scores
|
# Threshold all the scores
|
||||||
score_indices = sigmoid(colum_outputs[:, _YOLO_V2_SCORE])
|
score_indices = sigmoid(colum_outputs[:, _YOLO_V2_SCORE])
|
||||||
score_indices = np.nonzero(score_indices > self.score_threshold)
|
score_indices = np.nonzero(score_indices > self.threshold)
|
||||||
if isinstance(score_indices, tuple):
|
if isinstance(score_indices, tuple):
|
||||||
score_indices = score_indices[0]
|
score_indices = score_indices[0]
|
||||||
if not len(score_indices):
|
if not len(score_indices):
|
||||||
return []
|
return _NO_DETECTION
|
||||||
|
|
||||||
# Get the bounding boxes that have a valid score
|
# Get the bounding boxes that have a valid score
|
||||||
bb = np.take(colum_outputs, score_indices, axis=0)
|
bb = np.take(colum_outputs, score_indices, axis=0)
|
||||||
@ -159,7 +162,6 @@ class yolo_v5_postprocess:
|
|||||||
_YOLO_V5_CH = const(3)
|
_YOLO_V5_CH = const(3)
|
||||||
_YOLO_V5_SCORE = const(4)
|
_YOLO_V5_SCORE = const(4)
|
||||||
_YOLO_V5_CLASSES = const(5)
|
_YOLO_V5_CLASSES = const(5)
|
||||||
_NO_DETECTION = const(())
|
|
||||||
|
|
||||||
def __init__(self, threshold=0.6):
|
def __init__(self, threshold=0.6):
|
||||||
self.threshold = threshold
|
self.threshold = threshold
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user