mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
detector/plugins: added onnx cpu plugin
This commit is contained in:
parent
622a45ef12
commit
af3e40d2e1
@ -24,6 +24,7 @@ norfair == 2.2.*
|
|||||||
setproctitle == 1.3.*
|
setproctitle == 1.3.*
|
||||||
ws4py == 0.5.*
|
ws4py == 0.5.*
|
||||||
unidecode == 1.3.*
|
unidecode == 1.3.*
|
||||||
|
onnxruntime == 1.16.*
|
||||||
# Openvino Library - Custom built with MYRIAD support
|
# Openvino Library - Custom built with MYRIAD support
|
||||||
openvino @ https://github.com/NateMeyer/openvino-wheels/releases/download/multi-arch_2022.3.1/openvino-2022.3.1-1-cp39-cp39-manylinux_2_31_x86_64.whl; platform_machine == 'x86_64'
|
openvino @ https://github.com/NateMeyer/openvino-wheels/releases/download/multi-arch_2022.3.1/openvino-2022.3.1-1-cp39-cp39-manylinux_2_31_x86_64.whl; platform_machine == 'x86_64'
|
||||||
openvino @ https://github.com/NateMeyer/openvino-wheels/releases/download/multi-arch_2022.3.1/openvino-2022.3.1-1-cp39-cp39-linux_aarch64.whl; platform_machine == 'aarch64'
|
openvino @ https://github.com/NateMeyer/openvino-wheels/releases/download/multi-arch_2022.3.1/openvino-2022.3.1-1-cp39-cp39-linux_aarch64.whl; platform_machine == 'aarch64'
|
||||||
|
59
frigate/detectors/plugins/onnx.py
Normal file
59
frigate/detectors/plugins/onnx.py
Normal file
@ -0,0 +1,59 @@
|
|||||||
|
import logging
|
||||||
|
|
||||||
|
import sys
|
||||||
|
import os
|
||||||
|
import numpy as np
|
||||||
|
import ctypes
|
||||||
|
from pydantic import Field
|
||||||
|
from typing_extensions import Literal
|
||||||
|
import glob
|
||||||
|
import cv2
|
||||||
|
|
||||||
|
from frigate.detectors.detection_api import DetectionApi
|
||||||
|
from frigate.detectors.detector_config import BaseDetectorConfig
|
||||||
|
|
||||||
|
import frigate.detectors.yolo_utils as yolo_utils
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
DETECTOR_KEY = "onnx"
|
||||||
|
|
||||||
|
class ONNXDetectorConfig(BaseDetectorConfig):
|
||||||
|
type: Literal[DETECTOR_KEY]
|
||||||
|
|
||||||
|
class ONNXDetector(DetectionApi):
|
||||||
|
type_key = DETECTOR_KEY
|
||||||
|
|
||||||
|
def __init__(self, detector_config: ONNXDetectorConfig):
|
||||||
|
try:
|
||||||
|
import onnxruntime
|
||||||
|
|
||||||
|
logger.info(f"ONNX: loaded onnxruntime module")
|
||||||
|
except ModuleNotFoundError:
|
||||||
|
logger.error(
|
||||||
|
"ONNX: module loading failed, need 'pip install onnxruntime'?!?"
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
assert detector_config.model.model_type == 'yolov8', "ONNX: detector_config.model.model_type: only yolov8 supported"
|
||||||
|
assert detector_config.model.input_tensor == 'nhwc', "ONNX: detector_config.model.input_tensor: only nhwc supported"
|
||||||
|
if detector_config.model.input_pixel_format != 'rgb':
|
||||||
|
logger.warn("ONNX: detector_config.model.input_pixel_format: should be 'rgb' for yolov8, but '{detector_config.model.input_pixel_format}' specified!")
|
||||||
|
|
||||||
|
assert detector_config.model.path is not None, "ONNX: no model.path configured, please configure model.path and model.labelmap_path; some suggestions: " + ', '.join(glob.glob("/*.onnx")) + " and " + ', '.join(glob.glob("/*_labels.txt"))
|
||||||
|
|
||||||
|
path = detector_config.model.path
|
||||||
|
logger.info(f"ONNX: loading {detector_config.model.path}")
|
||||||
|
self.model = onnxruntime.InferenceSession(path)
|
||||||
|
logger.info(f"ONNX: {path} loaded")
|
||||||
|
|
||||||
|
def detect_raw(self, tensor_input):
|
||||||
|
model_input_name = self.model.get_inputs()[0].name
|
||||||
|
model_input_shape = self.model.get_inputs()[0].shape
|
||||||
|
|
||||||
|
tensor_input = yolo_utils.yolov8_preprocess(tensor_input, model_input_shape)
|
||||||
|
|
||||||
|
tensor_output = self.model.run(None, {model_input_name: tensor_input})[0]
|
||||||
|
|
||||||
|
return yolo_utils.yolov8_postprocess(model_input_shape, tensor_output)
|
||||||
|
|
Loading…
Reference in New Issue
Block a user