From de02ba761b634c9c666c4f81d7b8920311ae01c6 Mon Sep 17 00:00:00 2001 From: Indrek Mandre Date: Sun, 28 Jan 2024 11:14:37 +0200 Subject: [PATCH] detectors: simplified/optimized yolov8_postprocess --- frigate/detectors/yolo_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frigate/detectors/yolo_utils.py b/frigate/detectors/yolo_utils.py index 38083e705..43445b79f 100644 --- a/frigate/detectors/yolo_utils.py +++ b/frigate/detectors/yolo_utils.py @@ -18,7 +18,7 @@ def yolov8_postprocess(model_input_shape, tensor_output, box_count = 20): model_class_count = tensor_output.shape[1] - 4 probs = tensor_output[0, 4:, :] all_ids = np.argmax(probs, axis=0) - all_confidences = np.take(probs.T, model_class_count*np.arange(0, model_box_count) + all_ids) + all_confidences = probs.T[np.arange(model_box_count), all_ids] all_boxes = tensor_output[0, 0:4, :].T mask = (all_confidences > 0.30) class_ids = all_ids[mask]