administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-15 07:35:27 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Include device config

Browse Source
This commit is contained in:
Nicolas Mowen 2024-09-20 13:57:42 -06:00
parent 94a9196fc6
commit 80b182eff5
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
frigate/detectors/plugins/onnx.py
View File

@ -2,6 +2,7 @@ import logging
import os import os
import numpy as np import numpy as np
from pydantic import Field
from typing_extensions import Literal from typing_extensions import Literal
from frigate.detectors.detection_api import DetectionApi from frigate.detectors.detection_api import DetectionApi
@ -17,6 +18,7 @@ DETECTOR_KEY = "onnx"
class ONNXDetectorConfig(BaseDetectorConfig): class ONNXDetectorConfig(BaseDetectorConfig):
type: Literal[DETECTOR_KEY] type: Literal[DETECTOR_KEY]
device: str = Field(default="AUTO", title="Device Type")
class ONNXDetector(DetectionApi): class ONNXDetector(DetectionApi):
@ -36,7 +38,7 @@ class ONNXDetector(DetectionApi):
path = detector_config.model.path path = detector_config.model.path
logger.info(f"ONNX: loading {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 = [] options = []
for provider in providers: for provider in providers:
@ -57,7 +59,7 @@ class ONNXDetector(DetectionApi):
options.append( options.append(
{ {
"cache_dir": "/config/model_cache/openvino/ort", "cache_dir": "/config/model_cache/openvino/ort",
"device_type": "AUTO", "device_type": detector_config.device,
} }
) )
else: else: