From 2c5c8df6e9f1fb0c7ad108cc6949c2436be02aff Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Sun, 23 Mar 2025 13:11:39 -0600 Subject: [PATCH] Add support for openvino to run rf-detr models --- frigate/detectors/plugins/openvino.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py index 75d956500..9ef09d002 100644 --- a/frigate/detectors/plugins/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -10,7 +10,11 @@ from typing_extensions import Literal from frigate.const import MODEL_CACHE_DIR from frigate.detectors.detection_api import DetectionApi from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum -from frigate.util.model import post_process_dfine, post_process_yolov9 +from frigate.util.model import ( + post_process_dfine, + post_process_rfdetr, + post_process_yolov9, +) logger = logging.getLogger(__name__) @@ -25,11 +29,12 @@ class OvDetectorConfig(BaseDetectorConfig): class OvDetector(DetectionApi): type_key = DETECTOR_KEY supported_models = [ + ModelTypeEnum.dfine, + ModelTypeEnum.rfdetr, ModelTypeEnum.ssd, ModelTypeEnum.yolonas, ModelTypeEnum.yolov9, ModelTypeEnum.yolox, - ModelTypeEnum.dfine, ] def __init__(self, detector_config: OvDetectorConfig): @@ -185,6 +190,13 @@ class OvDetector(DetectionApi): if self.model_invalid: return detections + elif self.ov_model_type == ModelTypeEnum.rfdetr: + return post_process_rfdetr( + [ + infer_request.get_output_tensor(0).data, + infer_request.get_output_tensor(1).data, + ] + ) elif self.ov_model_type == ModelTypeEnum.ssd: results = infer_request.get_output_tensor(0).data[0][0]