rockchip: update dependencies and add script for model conversion (#15699)

* rockchip: update dependencies and add script for model conversion

* rockchip: update docs

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>

docker/rockchip/COCO/coco_subset_20.txt

@@ -0,0 +1,20 @@
+./subset/000000005001.jpg
+./subset/000000038829.jpg
+./subset/000000052891.jpg
+./subset/000000075612.jpg
+./subset/000000098261.jpg
+./subset/000000181542.jpg
+./subset/000000215245.jpg
+./subset/000000277005.jpg
+./subset/000000288685.jpg
+./subset/000000301421.jpg
+./subset/000000334371.jpg
+./subset/000000348481.jpg
+./subset/000000373353.jpg
+./subset/000000397681.jpg
+./subset/000000414673.jpg
+./subset/000000419312.jpg
+./subset/000000465822.jpg
+./subset/000000475732.jpg
+./subset/000000559707.jpg
+./subset/000000574315.jpg

docker/rockchip/Dockerfile

@@ -7,19 +7,23 @@ FROM wheels as rk-wheels
 COPY docker/main/requirements-wheels.txt /requirements-wheels.txt
 COPY docker/rockchip/requirements-wheels-rk.txt /requirements-wheels-rk.txt
 RUN sed -i "/https:\/\//d" /requirements-wheels.txt
+RUN sed -i "/onnxruntime/d" /requirements-wheels.txt
 RUN python3 -m pip config set global.break-system-packages true
 RUN pip3 wheel --wheel-dir=/rk-wheels -c /requirements-wheels.txt -r /requirements-wheels-rk.txt
+RUN rm -rf /rk-wheels/opencv_python-*
 
 FROM deps AS rk-frigate
 ARG TARGETARCH
 
 RUN --mount=type=bind,from=rk-wheels,source=/rk-wheels,target=/deps/rk-wheels \
-    pip3 install -U /deps/rk-wheels/*.whl --break-system-packages
+    pip3 install --no-deps -U /deps/rk-wheels/*.whl --break-system-packages
 
 WORKDIR /opt/frigate/
 COPY --from=rootfs / /
+COPY docker/rockchip/COCO /COCO
+COPY docker/rockchip/conv2rknn.py /opt/conv2rknn.py
 
-ADD https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.0.0/librknnrt.so /usr/lib/
+ADD https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.3.0/librknnrt.so /usr/lib/
 
 RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffmpeg
 RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffprobe

docker/rockchip/conv2rknn.py

@@ -0,0 +1,82 @@
+import os
+
+import rknn
+import yaml
+from rknn.api import RKNN
+
+try:
+    with open(rknn.__path__[0] + "/VERSION") as file:
+        tk_version = file.read().strip()
+except FileNotFoundError:
+    pass
+
+try:
+    with open("/config/conv2rknn.yaml", "r") as config_file:
+        configuration = yaml.safe_load(config_file)
+except FileNotFoundError:
+    raise Exception("Please place a config.yaml file in /config/conv2rknn.yaml")
+
+if configuration["config"] != None:
+    rknn_config = configuration["config"]
+else:
+    rknn_config = {}
+
+if not os.path.isdir("/config/model_cache/rknn_cache/onnx"):
+    raise Exception(
+        "Place the onnx models you want to convert to rknn format in /config/model_cache/rknn_cache/onnx"
+    )
+
+if "soc" not in configuration:
+    try:
+        with open("/proc/device-tree/compatible") as file:
+            soc = file.read().split(",")[-1].strip("\x00")
+    except FileNotFoundError:
+        raise Exception("Make sure to run docker in privileged mode.")
+
+    configuration["soc"] = [
+        soc,
+    ]
+
+if "quantization" not in configuration:
+    configuration["quantization"] = False
+
+if "output_name" not in configuration:
+    configuration["output_name"] = "{{input_basename}}"
+
+for input_filename in os.listdir("/config/model_cache/rknn_cache/onnx"):
+    for soc in configuration["soc"]:
+        quant = "i8" if configuration["quantization"] else "fp16"
+
+        input_path = "/config/model_cache/rknn_cache/onnx/" + input_filename
+        input_basename = input_filename[: input_filename.rfind(".")]
+
+        output_filename = (
+            configuration["output_name"].format(
+                quant=quant,
+                input_basename=input_basename,
+                soc=soc,
+                tk_version=tk_version,
+            )
+            + ".rknn"
+        )
+        output_path = "/config/model_cache/rknn_cache/" + output_filename
+
+        rknn_config["target_platform"] = soc
+
+        rknn = RKNN(verbose=True)
+        rknn.config(**rknn_config)
+
+        if rknn.load_onnx(model=input_path) != 0:
+            raise Exception("Error loading model.")
+
+        if (
+            rknn.build(
+                do_quantization=configuration["quantization"],
+                dataset="/COCO/coco_subset_20.txt",
+            )
+            != 0
+        ):
+            raise Exception("Error building model.")
+
+        if rknn.export_rknn(output_path) != 0:
+            raise Exception("Error exporting rknn model.")

docker/rockchip/requirements-wheels-rk.txt

@@ -1 +1,2 @@
-rknn-toolkit-lite2 @ https://github.com/MarcA711/rknn-toolkit2/releases/download/v2.0.0/rknn_toolkit_lite2-2.0.0b0-cp311-cp311-linux_aarch64.whl
+rknn-toolkit2 == 2.3.0
+rknn-toolkit-lite2 == 2.3.0