From 80b182eff57478c7d7e1f53473077f767b4e0170 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Fri, 20 Sep 2024 13:57:42 -0600 Subject: [PATCH] Include device config --- frigate/detectors/plugins/onnx.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index a3f2a25c6..7e7e426e8 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -2,6 +2,7 @@ import logging import os import numpy as np +from pydantic import Field from typing_extensions import Literal from frigate.detectors.detection_api import DetectionApi @@ -17,6 +18,7 @@ DETECTOR_KEY = "onnx" class ONNXDetectorConfig(BaseDetectorConfig): type: Literal[DETECTOR_KEY] + device: str = Field(default="AUTO", title="Device Type") class ONNXDetector(DetectionApi): @@ -36,7 +38,7 @@ class ONNXDetector(DetectionApi): path = detector_config.model.path logger.info(f"ONNX: loading {detector_config.model.path}") - providers = ort.get_available_providers() + providers = ["CPUExecutionProvider"] if detector_config.device == "CPU" else ort.get_available_providers() options = [] for provider in providers: @@ -57,7 +59,7 @@ class ONNXDetector(DetectionApi): options.append( { "cache_dir": "/config/model_cache/openvino/ort", - "device_type": "AUTO", + "device_type": detector_config.device, } ) else: