detector/plugins: added onnx cpu plugin

This commit is contained in:
Indrek Mandre 2024-01-27 22:39:15 +02:00
parent 622a45ef12
commit af3e40d2e1
2 changed files with 60 additions and 0 deletions

View File

@ -24,6 +24,7 @@ norfair == 2.2.*
setproctitle == 1.3.*
ws4py == 0.5.*
unidecode == 1.3.*
onnxruntime == 1.16.*
# 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-linux_aarch64.whl; platform_machine == 'aarch64'

View 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)