mirror of
https://github.com/openmv/openmv.git
synced 2025-11-04 14:49:50 +08:00
scripts/libraries: Simplify YOLO post-processing using keepdims.
This commit is contained in:
parent
6d21326233
commit
979b4e8ee6
@ -74,14 +74,13 @@ class yolo_v2_postprocess:
|
|||||||
|
|
||||||
def __init__(self, threshold=0.6, anchors=None, nms_threshold=0.1, nms_sigma=0.1):
|
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:
|
self.anchors = anchors
|
||||||
self.anchors = anchors
|
if self.anchors is None:
|
||||||
else:
|
|
||||||
self.anchors = np.array([[0.98830, 3.36060],
|
self.anchors = np.array([[0.98830, 3.36060],
|
||||||
[2.11940, 5.37590],
|
[2.11940, 5.37590],
|
||||||
[3.05200, 9.13360],
|
[3.05200, 9.13360],
|
||||||
[5.55170, 9.30660],
|
[5.55170, 9.30660],
|
||||||
[9.72600, 11.1422]], dtype=np.float)
|
[9.72600, 11.1422]])
|
||||||
self.anchors_len = len(self.anchors)
|
self.anchors_len = len(self.anchors)
|
||||||
self.nms_threshold = nms_threshold
|
self.nms_threshold = nms_threshold
|
||||||
self.nms_sigma = nms_sigma
|
self.nms_sigma = nms_sigma
|
||||||
@ -97,8 +96,8 @@ class yolo_v2_postprocess:
|
|||||||
return a - (b * (a // b))
|
return a - (b * (a // b))
|
||||||
|
|
||||||
def softmax(x):
|
def softmax(x):
|
||||||
e_x = np.exp(x - np.max(x))
|
e_x = np.exp(x - np.max(x, axis=1, keepdims=True))
|
||||||
return e_x / np.sum(e_x)
|
return e_x / np.sum(e_x, axis=1, keepdims=True)
|
||||||
|
|
||||||
# Reshape the output to a 2D array
|
# Reshape the output to a 2D array
|
||||||
colum_outputs = outputs[0].reshape((oh * ow * self.anchors_len,
|
colum_outputs = outputs[0].reshape((oh * ow * self.anchors_len,
|
||||||
@ -121,18 +120,13 @@ class yolo_v2_postprocess:
|
|||||||
bb_anchors = mod(score_indices, self.anchors_len)
|
bb_anchors = mod(score_indices, self.anchors_len)
|
||||||
|
|
||||||
# Get the anchor box information
|
# Get the anchor box information
|
||||||
bb_a_array = [self.anchors[i] for i in bb_anchors.tolist()]
|
bb_a_array = np.take(self.anchors, bb_anchors, axis=0)
|
||||||
bb_a_array = np.array(bb_a_array)
|
|
||||||
|
|
||||||
# Get the score information
|
# Get the score information
|
||||||
bb_scores = sigmoid(bb[:, _YOLO_V2_SCORE])
|
bb_scores = sigmoid(bb[:, _YOLO_V2_SCORE])
|
||||||
|
|
||||||
# Get the class information
|
# Get the class information
|
||||||
bb_classes = []
|
bb_classes = np.argmax(softmax(bb[:, _YOLO_V2_CLASSES:]), axis=1)
|
||||||
for i in range(len(score_indices)):
|
|
||||||
s = softmax(bb[i, _YOLO_V2_CLASSES:])
|
|
||||||
bb_classes.append(np.argmax(s))
|
|
||||||
bb_classes = np.array(bb_classes, dtype=np.uint16)
|
|
||||||
|
|
||||||
# Compute the bounding box information
|
# Compute the bounding box information
|
||||||
x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow
|
x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow
|
||||||
@ -192,8 +186,7 @@ class yolo_v5_postprocess:
|
|||||||
bb_scores = bb[:, _YOLO_V5_SCORE]
|
bb_scores = bb[:, _YOLO_V5_SCORE]
|
||||||
|
|
||||||
# Get the class information
|
# Get the class information
|
||||||
bb_classes = [np.argmax(bb[x, _YOLO_V5_CLASSES:]) for x in range(bb.shape[0])]
|
bb_classes = np.argmax(bb[:, _YOLO_V5_CLASSES:], axis=1)
|
||||||
bb_classes = np.array(bb_classes, dtype=np.uint16)
|
|
||||||
|
|
||||||
# Compute the bounding box information
|
# Compute the bounding box information
|
||||||
x_center = bb[:, _YOLO_V5_CX]
|
x_center = bb[:, _YOLO_V5_CX]
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user