From a4950da1feaa2e7b8a6761ffd72e80a5ab7953ec Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Sun, 23 Mar 2025 13:36:36 -0600 Subject: [PATCH] Add support for yolo generic --- frigate/detectors/plugins/onnx.py | 5 ++++- frigate/detectors/plugins/openvino.py | 6 +++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index 2679185a9..a10447b48 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -97,7 +97,10 @@ class ONNXDetector(DetectionApi): x_max / self.w, ] return detections - elif self.onnx_model_type == ModelTypeEnum.yolov9: + elif ( + self.onnx_model_type == ModelTypeEnum.yolov9 + or self.onnx_model_type == ModelTypeEnum.yologeneric + ): predictions: np.ndarray = tensor_output[0] return post_process_yolov9(predictions, self.w, self.h) else: diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py index 9ef09d002..d90352772 100644 --- a/frigate/detectors/plugins/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -34,6 +34,7 @@ class OvDetector(DetectionApi): ModelTypeEnum.ssd, ModelTypeEnum.yolonas, ModelTypeEnum.yolov9, + ModelTypeEnum.yologeneric, ModelTypeEnum.yolox, ] @@ -231,7 +232,10 @@ class OvDetector(DetectionApi): x_max / self.w, ] return detections - elif self.ov_model_type == ModelTypeEnum.yolov9: + elif ( + self.ov_model_type == ModelTypeEnum.yolov9 + or self.ov_model_type == ModelTypeEnum.yologeneric + ): out_tensor = infer_request.get_output_tensor(0).data return post_process_yolov9(out_tensor, self.w, self.h) elif self.ov_model_type == ModelTypeEnum.yolox: