diff --git a/frigate/detectors/yolo_utils.py b/frigate/detectors/yolo_utils.py index 43445b79f..79586890e 100644 --- a/frigate/detectors/yolo_utils.py +++ b/frigate/detectors/yolo_utils.py @@ -6,16 +6,15 @@ import cv2 logger = logging.getLogger(__name__) def yolov8_preprocess(tensor_input, model_input_shape): - # tensor_input must be nhwc - assert tensor_input.shape[3] == 3 - if tuple(tensor_input.shape[1:3]) != tuple(model_input_shape[2:4]): - logger.warn(f"yolov8_preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!") - # cv2.dnn.blobFromImage is faster than numpying it - return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False) + # tensor_input must be nhwc + assert tensor_input.shape[3] == 3 + if tuple(tensor_input.shape[1:3]) != tuple(model_input_shape[2:4]): + logger.warn(f"yolov8_preprocess: tensor_input.shape {tensor_input.shape} and model_input_shape {model_input_shape} do not match!") + # cv2.dnn.blobFromImage is faster than numpying it + return cv2.dnn.blobFromImage(tensor_input[0], 1.0 / 255, (model_input_shape[3], model_input_shape[2]), None, swapRB=False) def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20): model_box_count = tensor_output.shape[2] - model_class_count = tensor_output.shape[1] - 4 probs = tensor_output[0, 4:, :] all_ids = np.argmax(probs, axis=0) all_confidences = probs.T[np.arange(model_box_count), all_ids]