diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index 7da7fe1bd..c63cee0a9 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -29,6 +29,8 @@ import math import image from ml.utils import NMS +from micropython import const +from ulab import numpy as np # FOMO generates an image per class, where each pixel represents the centroid @@ -57,3 +59,94 @@ class fomo_postprocess: ) nms.add_bounding_box(x, y, x + w, y + h, score, i) return nms.get_bounding_boxes() + + +class yolo_v2_postprocess: + _YOLO_V2_TX = const(0) + _YOLO_V2_TY = const(1) + _YOLO_V2_TW = const(2) + _YOLO_V2_TH = const(3) + _YOLO_V2_SCORE = const(4) + _YOLO_V2_CLASSES = const(5) + + def __init__(self, score_threshold=0.6, anchors=None): + self.score_threshold = score_threshold + if anchors is not None: + self.anchors = anchors + else: + 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) + self.anchors_len = len(self.anchors) + + def __call__(self, model, inputs, outputs): + ob, oh, ow, oc = model.output_shape[0] + class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES + + def sigmoid(x): + return 1.0 / (1.0 + np.exp(-x)) + + def mod(a, b): + return a - (b * (a // b)) + + def softmax(x): + e_x = np.exp(x - np.max(x)) + return e_x / np.sum(e_x) + + # Reshape the output to a 2D array + colum_outputs = outputs[0].reshape((oh * ow * self.anchors_len, + _YOLO_V2_CLASSES + class_count)) + + # Threshold all the scores + score_indices = sigmoid(colum_outputs[:, _YOLO_V2_SCORE]) + score_indices = np.nonzero(score_indices > self.score_threshold) + if isinstance(score_indices, tuple): + score_indices = score_indices[0] + if not len(score_indices): + return [] + + # Get the bounding boxes that have a valid score + bb = np.take(colum_outputs, score_indices, axis=0) + + # Extract rows, columns, and anchor indices + bb_rows = score_indices // (ow * self.anchors_len) + bb_cols = mod(score_indices // self.anchors_len, ow) + 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) + + # 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) + + # 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 + + # 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 + + nms = NMS(iw, ih, inputs[0].roi) + for i in range(len(bb)): + 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]) + return nms.get_bounding_boxes() diff --git a/src/omv/boards/ARDUINO_GIGA/ulab_config.h b/src/omv/boards/ARDUINO_GIGA/ulab_config.h index 00788f3ba..be1b0fae6 100644 --- a/src/omv/boards/ARDUINO_GIGA/ulab_config.h +++ b/src/omv/boards/ARDUINO_GIGA/ulab_config.h @@ -16,4 +16,5 @@ #define NDARRAY_BINARY_USES_FUN_POINTER (1) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/ARDUINO_NANO_33_BLE_SENSE/ulab_config.h b/src/omv/boards/ARDUINO_NANO_33_BLE_SENSE/ulab_config.h index e8fecc0aa..9015065b3 100644 --- a/src/omv/boards/ARDUINO_NANO_33_BLE_SENSE/ulab_config.h +++ b/src/omv/boards/ARDUINO_NANO_33_BLE_SENSE/ulab_config.h @@ -13,4 +13,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/ARDUINO_NANO_RP2040_CONNECT/ulab_config.h b/src/omv/boards/ARDUINO_NANO_RP2040_CONNECT/ulab_config.h index e8fecc0aa..9015065b3 100644 --- a/src/omv/boards/ARDUINO_NANO_RP2040_CONNECT/ulab_config.h +++ b/src/omv/boards/ARDUINO_NANO_RP2040_CONNECT/ulab_config.h @@ -13,4 +13,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/ARDUINO_NICLA_VISION/ulab_config.h b/src/omv/boards/ARDUINO_NICLA_VISION/ulab_config.h index ebdbed9e2..50a8a9949 100644 --- a/src/omv/boards/ARDUINO_NICLA_VISION/ulab_config.h +++ b/src/omv/boards/ARDUINO_NICLA_VISION/ulab_config.h @@ -15,4 +15,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) #define NDARRAY_BINARY_USES_FUN_POINTER (1) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/ARDUINO_PORTENTA_H7/ulab_config.h b/src/omv/boards/ARDUINO_PORTENTA_H7/ulab_config.h index d2093a723..f3aee6cf4 100644 --- a/src/omv/boards/ARDUINO_PORTENTA_H7/ulab_config.h +++ b/src/omv/boards/ARDUINO_PORTENTA_H7/ulab_config.h @@ -14,4 +14,5 @@ #define NDARRAY_BINARY_USES_FUN_POINTER (1) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/OPENMV3/ulab_config.h b/src/omv/boards/OPENMV3/ulab_config.h index d2093a723..f3aee6cf4 100644 --- a/src/omv/boards/OPENMV3/ulab_config.h +++ b/src/omv/boards/OPENMV3/ulab_config.h @@ -14,4 +14,5 @@ #define NDARRAY_BINARY_USES_FUN_POINTER (1) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/OPENMV4/ulab_config.h b/src/omv/boards/OPENMV4/ulab_config.h index e8fecc0aa..9015065b3 100644 --- a/src/omv/boards/OPENMV4/ulab_config.h +++ b/src/omv/boards/OPENMV4/ulab_config.h @@ -13,4 +13,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/OPENMV4P/ulab_config.h b/src/omv/boards/OPENMV4P/ulab_config.h index e8fecc0aa..9015065b3 100644 --- a/src/omv/boards/OPENMV4P/ulab_config.h +++ b/src/omv/boards/OPENMV4P/ulab_config.h @@ -13,4 +13,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/OPENMVPT/ulab_config.h b/src/omv/boards/OPENMVPT/ulab_config.h index e8fecc0aa..9015065b3 100644 --- a/src/omv/boards/OPENMVPT/ulab_config.h +++ b/src/omv/boards/OPENMVPT/ulab_config.h @@ -13,4 +13,5 @@ #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/OPENMV_RT1060/ulab_config.h b/src/omv/boards/OPENMV_RT1060/ulab_config.h index 4c70739e7..b0adee677 100644 --- a/src/omv/boards/OPENMV_RT1060/ulab_config.h +++ b/src/omv/boards/OPENMV_RT1060/ulab_config.h @@ -12,4 +12,5 @@ #define ULAB_MAX_DIMS (4) #define ULAB_SUPPORTS_COMPLEX (0) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (1) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/boards/RPI_PICO/ulab_config.h b/src/omv/boards/RPI_PICO/ulab_config.h index a7e8e7711..e538ecfe3 100644 --- a/src/omv/boards/RPI_PICO/ulab_config.h +++ b/src/omv/boards/RPI_PICO/ulab_config.h @@ -13,4 +13,5 @@ #define NDARRAY_BINARY_USES_FUN_POINTER (1) #define ULAB_SCIPY_HAS_OPTIMIZE_MODULE (0) #define ULAB_SCIPY_HAS_SPECIAL_MODULE (0) +#define ULAB_FFT_IS_NUMPY_COMPATIBLE (0) #endif //__ULAB_CONFIG_H__ diff --git a/src/omv/modules/ulab b/src/omv/modules/ulab index 65c941a80..2b74236c8 160000 --- a/src/omv/modules/ulab +++ b/src/omv/modules/ulab @@ -1 +1 @@ -Subproject commit 65c941a8059afe1cfd6f4c2b15d0ade798dc24f2 +Subproject commit 2b74236c8ce278d876e410c00e90de5bf5e9980e