From 35cc6a17492fd871f90582e73d8dda8ed4661c55 Mon Sep 17 00:00:00 2001 From: Indrek Mandre Date: Tue, 6 Feb 2024 16:12:37 +0200 Subject: [PATCH] add openvino/yolov8 support for label aggregation --- frigate/detectors/plugins/openvino.py | 26 ++++---------------------- 1 file changed, 4 insertions(+), 22 deletions(-) diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py index 5d748f1e7..fc2ac4682 100644 --- a/frigate/detectors/plugins/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -51,6 +51,8 @@ class OvDetector(DetectionApi): logger.info(f"YOLOX model has {self.num_classes} classes") self.set_strides_grids() + self.class_aggregation = yolo_utils.generate_class_aggregation_from_config(detector_config) + def set_strides_grids(self): grids = [] expanded_strides = [] @@ -135,28 +137,8 @@ class OvDetector(DetectionApi): ) return detections elif self.ov_model_type == ModelTypeEnum.yolov8: - out_tensor = infer_request.get_output_tensor() - results = out_tensor.data[0] - output_data = np.transpose(results) - scores = np.max(output_data[:, 4:], axis=1) - if len(scores) == 0: - return np.zeros((20, 6), np.float32) - scores = np.expand_dims(scores, axis=1) - # add scores to the last column - dets = np.concatenate((output_data, scores), axis=1) - # filter out lines with scores below threshold - dets = dets[dets[:, -1] > 0.5, :] - # limit to top 20 scores, descending order - ordered = dets[dets[:, -1].argsort()[::-1]][:20] - detections = np.zeros((20, 6), np.float32) - - for i, object_detected in enumerate(ordered): - detections[i] = self.process_yolo( - np.argmax(object_detected[4:-1]), - object_detected[-1], - object_detected[:4], - ) - return detections + out_tensor = infer_request.get_output_tensor().data + return yolo_utils.yolov8_postprocess(self.interpreter.inputs[0].shape, out_tensor, class_aggregation = self.class_aggregation) elif self.ov_model_type == ModelTypeEnum.yolov5: out_tensor = infer_request.get_output_tensor() output_data = out_tensor.data[0]