From f41b23d9a66b2ffcebcc194748dc322c9df34c91 Mon Sep 17 00:00:00 2001
From: MarcA711 <40744649+MarcA711@users.noreply.github.com>
Date: Fri, 17 May 2024 15:32:19 +0000
Subject: [PATCH] update rknn

---
 docker/rockchip/Dockerfile                 | 3 +--
 docker/rockchip/requirements-wheels-rk.txt | 3 +--
 frigate/detectors/plugins/rknn.py          | 8 ++++----
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/docker/rockchip/Dockerfile b/docker/rockchip/Dockerfile
index d42117f7c..d2404702f 100644
--- a/docker/rockchip/Dockerfile
+++ b/docker/rockchip/Dockerfile
@@ -18,8 +18,7 @@ RUN --mount=type=bind,from=rk-wheels,source=/rk-wheels,target=/deps/rk-wheels \
 WORKDIR /opt/frigate/
 COPY --from=rootfs / /
 
-ADD https://github.com/MarcA711/rknpu2/releases/download/v1.5.2/librknnrt_rk356x.so /usr/lib/
-ADD https://github.com/MarcA711/rknpu2/releases/download/v1.5.2/librknnrt_rk3588.so /usr/lib/
+ADD https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.0.0/librknnrt.so /usr/lib/
 
 # TODO removed models, other models support may need to be added back in
 
diff --git a/docker/rockchip/requirements-wheels-rk.txt b/docker/rockchip/requirements-wheels-rk.txt
index 9a3fe5c77..c56b69b66 100644
--- a/docker/rockchip/requirements-wheels-rk.txt
+++ b/docker/rockchip/requirements-wheels-rk.txt
@@ -1,2 +1 @@
-hide-warnings == 0.17
-rknn-toolkit-lite2 @ https://github.com/MarcA711/rknn-toolkit2/releases/download/v1.5.2/rknn_toolkit_lite2-1.5.2-cp39-cp39-linux_aarch64.whl
\ No newline at end of file
+rknn-toolkit-lite2 @ https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.0.0/rknn_toolkit_lite2-2.0.0b0-cp39-cp39-linux_aarch64.whl
\ No newline at end of file
diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py
index 9f6d3a61c..5e830880c 100644
--- a/frigate/detectors/plugins/rknn.py
+++ b/frigate/detectors/plugins/rknn.py
@@ -32,6 +32,7 @@ class Rknn(DetectionApi):
     type_key = DETECTOR_KEY
 
     def __init__(self, config: RknnDetectorConfig):
+        super().__init__(config)
         self.height = config.model.height
         self.width = config.model.width
         core_mask = 2**config.num_cores - 1
@@ -108,6 +109,8 @@ class Rknn(DetectionApi):
                     f"Model {model_path} is unsupported. Provide your own model or choose one of the following: {supported_models_str}"
                 )
 
+        return model_props
+
     def download_model(self, filename):
         if not os.path.isdir(model_chache_dir):
             os.mkdir(model_chache_dir)
@@ -133,13 +136,10 @@ class Rknn(DetectionApi):
                 'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
             )
 
-    def post_process(self, output):
-        pass
-
     def detect_raw(self, tensor_input):
         output = self.rknn.inference(
             [
                 tensor_input,
             ]
         )
-        return self.postprocess(output[0])
+        return self.post_process(output)