From 979b4e8ee6fe34d6685c68a7675d10067c90b04d Mon Sep 17 00:00:00 2001 From: "Kwabena W. Agyeman" Date: Wed, 22 Jan 2025 17:20:04 -0800 Subject: [PATCH] scripts/libraries: Simplify YOLO post-processing using keepdims. --- scripts/libraries/ml/ml/postprocessing.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index d08c0720e..cd6b833fa 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -74,14 +74,13 @@ class yolo_v2_postprocess: 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 - else: + self.anchors = anchors + if self.anchors is None: self.anchors = np.array([[0.98830, 3.36060], [2.11940, 5.37590], [3.05200, 9.13360], [5.55170, 9.30660], - [9.72600, 11.1422]], dtype=np.float) + [9.72600, 11.1422]]) self.anchors_len = len(self.anchors) self.nms_threshold = nms_threshold self.nms_sigma = nms_sigma @@ -97,8 +96,8 @@ class yolo_v2_postprocess: return a - (b * (a // b)) def softmax(x): - e_x = np.exp(x - np.max(x)) - return e_x / np.sum(e_x) + e_x = np.exp(x - np.max(x, axis=1, keepdims=True)) + return e_x / np.sum(e_x, axis=1, keepdims=True) # Reshape the output to a 2D array 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) # Get the anchor box information - bb_a_array = [self.anchors[i] for i in bb_anchors.tolist()] - bb_a_array = np.array(bb_a_array) + bb_a_array = np.take(self.anchors, bb_anchors, axis=0) # Get the score information bb_scores = sigmoid(bb[:, _YOLO_V2_SCORE]) # Get the class information - bb_classes = [] - 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) + bb_classes = np.argmax(softmax(bb[:, _YOLO_V2_CLASSES:]), axis=1) # Compute the bounding box information x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow @@ -192,8 +186,7 @@ class yolo_v5_postprocess: bb_scores = bb[:, _YOLO_V5_SCORE] # Get the class information - bb_classes = [np.argmax(bb[x, _YOLO_V5_CLASSES:]) for x in range(bb.shape[0])] - bb_classes = np.array(bb_classes, dtype=np.uint16) + bb_classes = np.argmax(bb[:, _YOLO_V5_CLASSES:], axis=1) # Compute the bounding box information x_center = bb[:, _YOLO_V5_CX]