Add support for RKNN

docker/main/Dockerfile

@@ -90,6 +90,10 @@ RUN /bin/mkdir -p '/usr/local/lib' && \
     /usr/bin/install -c -m 644 libusb-1.0.pc '/usr/local/lib/pkgconfig' && \
     ldconfig
 
+# Download libs needed for RKNN
+FROM wget AS rknn-deps
+RUN wget -P /rknn-lib https://github.com/rockchip-linux/rknpu2/raw/master/runtime/RK3588/Linux/librknn_api/aarch64/librknnrt.so
+
 FROM wget AS models
 
 # Get model and labels
@@ -159,6 +163,7 @@ FROM scratch AS deps-rootfs
 COPY --from=nginx /usr/local/nginx/ /usr/local/nginx/
 COPY --from=go2rtc /rootfs/ /
 COPY --from=libusb-build /usr/local/lib /usr/local/lib
+COPY --from=rknn-deps /rknn-lib /usr/lib
 COPY --from=s6-overlay /rootfs/ /
 COPY --from=models /rootfs/ /
 COPY docker/main/rootfs/ /

docker/main/install_deps.sh

@@ -72,6 +72,12 @@ if [[ "${TARGETARCH}" == "arm64" ]]; then
         libva-drm2 mesa-va-drivers
 fi
 
+# Needed for RKNN
+if [[ "${TARGETARCH}" == "arm64" ]]; then
+    apt-get -qq install --no-install-recommends --no-install-suggests -y \
+        gcc
+fi
+
 apt-get purge gnupg apt-transport-https xz-utils -y
 apt-get clean autoclean -y
 apt-get autoremove --purge -y

docker/main/requirements-wheels.txt

@@ -26,3 +26,6 @@ ws4py == 0.5.*
 # 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'
+# RKNN-Toolkit-Lite2
+hide-warnings == 0.17 # required to supress warnings from RKNNLite
+rknn-toolkit-lite2 @ https://github.com/rockchip-linux/rknn-toolkit2/raw/master/rknn_toolkit_lite2/packages/rknn_toolkit_lite2-1.5.2-cp39-cp39-linux_aarch64.whl; platform_machine == 'aarch64'

frigate/detectors/plugins/rknn.py

@@ -0,0 +1,96 @@
+import logging
+import numpy as np
+import cv2
+import cv2.dnn
+from hide_warnings import hide_warnings
+
+from frigate.detectors.detection_api import DetectionApi
+from frigate.detectors.detector_config import BaseDetectorConfig
+from typing import Literal
+from pydantic import Field
+
+logger = logging.getLogger(__name__)
+
+DETECTOR_KEY = "rknn"
+
+class RknnDetectorConfig(BaseDetectorConfig):
+    type: Literal[DETECTOR_KEY]
+    core_mask: int = Field(default=0, ge=0, le= 7, title="Core mask for RKNN.")
+    score_thresh: float = Field(default=0.5, ge=0, le= 1, title="Minimal confidence for detection.")
+    nms_thresh: float = Field(default=0.45, ge=0, le= 1, title="IoU threshold for non-maximum suppression.")
+
+class Rknn(DetectionApi):
+    type_key = DETECTOR_KEY
+
+    def __init__(self, config: RknnDetectorConfig):
+        self.height = config.model.height
+        self.width = config.model.width
+        self.score_thresh = config.score_thresh
+        self.nms_thresh = config.nms_thresh
+
+        self.model_path = config.model.path or "/models/rknn/yolov8m-320x320.rknn"
+        self.core_mask = config.core_mask
+
+        from rknnlite.api import RKNNLite
+        self.rknn = RKNNLite(verbose=False)
+        if self.rknn.load_rknn(self.model_path) != 0:
+            logger.error('Error initializing rknn model.')
+        if self.rknn.init_runtime(core_mask=self.core_mask) != 0:
+            logger.error('Error initializing rknn runtime.')
+
+    def __del__(self):
+        self.rknn.release()
+
+    def postprocess(self, results):
+        """
+        Processes yolov8 output.
+
+        Args:
+        results: array with shape: (1, 84, n, 1) where n depends on yolov8 model size (for 320x320 model n=2100)
+
+        Returns:
+        detections: array with shape (20, 6) with 20 rows of (class, confidence, y_min, x_min, y_max, x_max)
+        """
+
+        results = np.transpose(results[0,:,:,0]) # array shape (2100, 84)
+        classes = np.argmax(results[:, 4:], axis=1) # array shape (2100,); index of class with max confidence of each row
+        scores = np.max(results[:, 4:], axis=1) # array shape (2100,); max confidence of each row
+
+        # array shape (2100, 4); bounding box of each row
+        boxes = np.transpose(
+            np.vstack((
+                results[:,0] - 0.5 * results[:,2],
+                results[:,1] - 0.5 * results[:,3],
+                results[:,2],
+                results[:,3])
+            )
+        )
+
+        # indices of rows with confidence > SCORE_THRESH with Non-maximum Suppression (NMS)
+        result_boxes = cv2.dnn.NMSBoxes(boxes, scores, self.score_thresh, self.nms_thresh, 0.5)
+
+        detections = np.zeros((20, 6), np.float32)
+
+        for i in range(len(result_boxes)):
+            if i >= 20:
+                break
+
+            index = result_boxes[i]
+            detections[i] = [
+                classes[index],
+                scores[index],
+                (boxes[index][1]) / self.height,
+                (boxes[index][0]) / self.width,
+                (boxes[index][1] + boxes[index][3]) / self.height,
+                (boxes[index][0] + boxes[index][2]) / self.width
+            ]
+
+        return detections
+
+    @hide_warnings
+    def inference(self, tensor_input):
+        return self.rknn.inference(inputs=tensor_input)
+
+    def detect_raw(self, tensor_input):
+        output = self.inference([tensor_input,])
+        return self.postprocess(output[0])