diff --git a/docker/main/Dockerfile b/docker/main/Dockerfile index aee067b25..9939b67f2 100644 --- a/docker/main/Dockerfile +++ b/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/ / diff --git a/docker/main/install_deps.sh b/docker/main/install_deps.sh index 43fff479b..666d595bc 100755 --- a/docker/main/install_deps.sh +++ b/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 diff --git a/docker/main/requirements-wheels.txt b/docker/main/requirements-wheels.txt index 37028d33b..093a51901 100644 --- a/docker/main/requirements-wheels.txt +++ b/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' diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py new file mode 100644 index 000000000..9b8b590bb --- /dev/null +++ b/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])