diff --git a/scripts/libraries/ml/ml/postprocessing.py b/scripts/libraries/ml/ml/postprocessing.py index cf32f55db..8254dcdd1 100644 --- a/scripts/libraries/ml/ml/postprocessing.py +++ b/scripts/libraries/ml/ml/postprocessing.py @@ -38,6 +38,23 @@ def mod(a, b): return a - (b * (a // b)) +def threshold(scores, threshold, scale, find_max=False, find_max_axis=1): + if scale > 0: + if find_max: + scores = np.max(scores, axis=find_max_axis) + return np.nonzero(scores > threshold)[0] + else: + if find_max: + scores = np.min(scores, axis=find_max_axis) + return np.nonzero(scores < threshold)[0] + + +def quantize(model, value): + if model.output_dtype[0] == 'f': + return value + return (value / model.output_scale[0]) + model.output_zero_point[0] + + def dequantize(model, value): if model.output_dtype[0] == 'f': return value @@ -57,13 +74,15 @@ class fomo_postprocess: def __call__(self, model, inputs, outputs): ob, oh, ow, oc = model.output_shape[0] + scale = model.output_scale[0] + t = quantize(model, self.threshold) # Reshape the output to a 2D array row_outputs = outputs[0].reshape((oh * ow, oc)) # Threshold all the scores - score_indices = np.max(dequantize(model, row_outputs[:, _FOMO_CLASSES:]), axis=1) - score_indices = np.nonzero(score_indices > self.threshold)[0] + score_indices = row_outputs[:, _FOMO_CLASSES:] + score_indices = threshold(score_indices, t, scale, find_max=True, find_max_axis=1) if not len(score_indices): return _NO_DETECTION @@ -121,8 +140,9 @@ class yolo_v2_postprocess: self.nms_sigma = nms_sigma def __call__(self, model, inputs, outputs): - ob, oh, ow, oc = model.output_shape[0] - class_count = (oc // self.anchors_len) - _YOLO_V2_CLASSES + + def logit(x): + return np.log(x / (1.0 - x)) def sigmoid(x): return 1.0 / (1.0 + np.exp(-x)) @@ -131,13 +151,18 @@ class yolo_v2_postprocess: e_x = np.exp(x - np.max(x, axis=1, keepdims=True)) return e_x / np.sum(e_x, axis=1, keepdims=True) + ob, oh, ow, oc = model.output_shape[0] + scale = model.output_scale[0] + t = quantize(model, logit(self.threshold)) + 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)) # Threshold all the scores - score_indices = sigmoid(dequantize(model, row_outputs[:, _YOLO_V2_SCORE])) - score_indices = np.nonzero(score_indices > self.threshold)[0] + score_indices = row_outputs[:, _YOLO_V2_SCORE] + score_indices = threshold(score_indices, t, scale) if not len(score_indices): return _NO_DETECTION @@ -209,14 +234,16 @@ class yolo_v5_postprocess: def __call__(self, model, inputs, outputs): oh, ow, oc = model.output_shape[0] + scale = model.output_scale[0] + t = quantize(model, self.threshold) class_count = oc - _YOLO_V5_CLASSES # Reshape the output to a 2D array row_outputs = outputs[0].reshape((oh * ow, _YOLO_V5_CLASSES + class_count)) # Threshold all the scores - score_indices = dequantize(model, row_outputs[:, _YOLO_V5_SCORE]) - score_indices = np.nonzero(score_indices > self.threshold)[0] + score_indices = row_outputs[:, _YOLO_V5_SCORE] + score_indices = threshold(score_indices, t, scale) if not len(score_indices): return _NO_DETECTION @@ -263,14 +290,16 @@ class yolo_v8_postprocess: def __call__(self, model, inputs, outputs): oh, ow, oc = model.output_shape[0] + scale = model.output_scale[0] + t = quantize(model, self.threshold) class_count = ow - _YOLO_V8_CLASSES # Reshape the output to a 2D array column_outputs = outputs[0].reshape((oh * (_YOLO_V8_CLASSES + class_count), oc)) # Threshold all the scores - score_indices = np.max(dequantize(model, column_outputs[_YOLO_V8_CLASSES:, :]), axis=0) - score_indices = np.nonzero(score_indices > self.threshold)[0] + score_indices = column_outputs[_YOLO_V8_CLASSES:, :] + score_indices = threshold(score_indices, t, scale, find_max=True, find_max_axis=0) if not len(score_indices): return _NO_DETECTION