diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index 07311b56b..763ea1b64 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -107,20 +107,23 @@ class fomo_postprocess: # Get the class information bb_classes = np.argmax(bb[:, _FOMO_CLASSES:], axis=1) + _FOMO_CLASSES + # Compute the bounding box information + x_center = (bb_cols + 0.5) / ow + y_center = (bb_rows + 0.5) / oh + w_rel = np.full(len(bb_cols), self.w_scale / ow) * 0.5 + h_rel = np.full(len(bb_rows), self.h_scale / oh) * 0.5 + # Scale the bounding boxes to have enough integer precision for NMS ib, ih, iw, ic = model.input_shape[0] - x_center = ((bb_cols + 0.5) / ow) * iw - y_center = ((bb_rows + 0.5) / oh) * ih - w_rel = np.full(len(bb_cols), self.w_scale / ow) * iw - h_rel = np.full(len(bb_rows), self.h_scale / oh) * ih + xmin = (x_center - w_rel) * iw + ymin = (y_center - h_rel) * ih + xmax = (x_center + w_rel) * iw + ymax = (y_center + h_rel) * ih nms = NMS(iw, ih, inputs[0].roi) for i in range(bb.shape[0]): - nms.add_bounding_box(x_center[i] - (w_rel[i] / 2), - y_center[i] - (h_rel[i] / 2), - x_center[i] + (w_rel[i] / 2), - y_center[i] + (h_rel[i] / 2), - bb_scores[i], bb_classes[i]) + nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], bb_scores[i], bb_classes[i]) + return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma) @@ -159,8 +162,7 @@ class yolo_v2_postprocess: class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES # Reshape the output to a 2D array - row_outputs = outputs[0].reshape((oh * ow * self.anchors_len, - _YOLO_V2_CLASSES + class_count)) + row_outputs = outputs[0].reshape((oh * ow * self.anchors_len, _YOLO_V2_CLASSES + class_count)) # Threshold all the scores score_indices = row_outputs[:, _YOLO_V2_SCORE] @@ -188,23 +190,20 @@ class yolo_v2_postprocess: # Compute the bounding box information x_center = (bb_cols + sigmoid(bb[:, _YOLO_V2_TX])) / ow y_center = (bb_rows + sigmoid(bb[:, _YOLO_V2_TY])) / oh - w_rel = (bb_a_array[:, 0] * np.exp(bb[:, _YOLO_V2_TW])) / ow - h_rel = (bb_a_array[:, 1] * np.exp(bb[:, _YOLO_V2_TH])) / oh + w_rel = ((bb_a_array[:, 0] * np.exp(bb[:, _YOLO_V2_TW])) / ow) * 0.5 + h_rel = ((bb_a_array[:, 1] * np.exp(bb[:, _YOLO_V2_TH])) / oh) * 0.5 # Scale the bounding boxes to have enough integer precision for NMS ib, ih, iw, ic = model.input_shape[0] - x_center = x_center * iw - y_center = y_center * ih - w_rel = w_rel * iw - h_rel = h_rel * ih + xmin = (x_center - w_rel) * iw + ymin = (y_center - h_rel) * ih + xmax = (x_center + w_rel) * iw + ymax = (y_center + h_rel) * ih nms = NMS(iw, ih, inputs[0].roi) for i in range(bb.shape[0]): - nms.add_bounding_box(x_center[i] - (w_rel[i] / 2), - y_center[i] - (h_rel[i] / 2), - x_center[i] + (w_rel[i] / 2), - y_center[i] + (h_rel[i] / 2), - bb_scores[i], bb_classes[i]) + nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], bb_scores[i], bb_classes[i]) + return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma) @@ -273,8 +272,8 @@ class yolo_v5_postprocess: nms = NMS(iw, ih, inputs[0].roi) for i in range(bb.shape[0]): - nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], - bb_scores[i], bb_classes[i]) + nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], bb_scores[i], bb_classes[i]) + return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma) @@ -329,8 +328,8 @@ class yolo_v8_postprocess: nms = NMS(iw, ih, inputs[0].roi) for i in range(bb.shape[0]): - nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], - bb_scores[i], bb_classes[i]) + nms.add_bounding_box(xmin[i], ymin[i], xmax[i], ymax[i], bb_scores[i], bb_classes[i]) + return nms.get_bounding_boxes(threshold=self.nms_threshold, sigma=self.nms_sigma)