Add support for openvino to run rf-detr models

This commit is contained in:
Nicolas Mowen 2025-03-23 13:11:39 -06:00
parent b7fcd41737
commit 2c5c8df6e9

View File

@ -10,7 +10,11 @@ from typing_extensions import Literal
from frigate.const import MODEL_CACHE_DIR from frigate.const import MODEL_CACHE_DIR
from frigate.detectors.detection_api import DetectionApi from frigate.detectors.detection_api import DetectionApi
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum 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__) logger = logging.getLogger(__name__)
@ -25,11 +29,12 @@ class OvDetectorConfig(BaseDetectorConfig):
class OvDetector(DetectionApi): class OvDetector(DetectionApi):
type_key = DETECTOR_KEY type_key = DETECTOR_KEY
supported_models = [ supported_models = [
ModelTypeEnum.dfine,
ModelTypeEnum.rfdetr,
ModelTypeEnum.ssd, ModelTypeEnum.ssd,
ModelTypeEnum.yolonas, ModelTypeEnum.yolonas,
ModelTypeEnum.yolov9, ModelTypeEnum.yolov9,
ModelTypeEnum.yolox, ModelTypeEnum.yolox,
ModelTypeEnum.dfine,
] ]
def __init__(self, detector_config: OvDetectorConfig): def __init__(self, detector_config: OvDetectorConfig):
@ -185,6 +190,13 @@ class OvDetector(DetectionApi):
if self.model_invalid: if self.model_invalid:
return detections 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: elif self.ov_model_type == ModelTypeEnum.ssd:
results = infer_request.get_output_tensor(0).data[0][0] results = infer_request.get_output_tensor(0).data[0][0]