From af3e40d2e14c83e2b2cd5f053a4c5d6fd9a03839 Mon Sep 17 00:00:00 2001
From: Indrek Mandre <indrek@mare.ee>
Date: Sat, 27 Jan 2024 22:39:15 +0200
Subject: [PATCH] detector/plugins: added onnx cpu plugin

---
 docker/main/requirements-wheels.txt |  1 +
 frigate/detectors/plugins/onnx.py   | 59 +++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)
 create mode 100644 frigate/detectors/plugins/onnx.py

diff --git a/docker/main/requirements-wheels.txt b/docker/main/requirements-wheels.txt
index f4167744e..cb5f261e6 100644
--- a/docker/main/requirements-wheels.txt
+++ b/docker/main/requirements-wheels.txt
@@ -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'
diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py
new file mode 100644
index 000000000..c98665d4d
--- /dev/null
+++ b/frigate/detectors/plugins/onnx.py
@@ -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)
+