mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-05-05 13:07:44 +03:00
Add support for openvino to run rf-detr models
This commit is contained in:
parent
b7fcd41737
commit
2c5c8df6e9
@ -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]
|
||||||
|
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user