From de6ae9adb7943c8965365c099de4fd6029d5e939 Mon Sep 17 00:00:00 2001 From: MarcA711 <40744649+MarcA711@users.noreply.github.com> Date: Mon, 20 May 2024 10:58:43 +0000 Subject: [PATCH] add yolonas models --- docker/rockchip/Dockerfile | 6 +++++- frigate/detectors/plugins/rknn.py | 8 ++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docker/rockchip/Dockerfile b/docker/rockchip/Dockerfile index d2404702f..35b258070 100644 --- a/docker/rockchip/Dockerfile +++ b/docker/rockchip/Dockerfile @@ -20,7 +20,11 @@ COPY --from=rootfs / / 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 +ADD https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/deci-fp16-yolonas_s-rk3562-v2.0.0-1.rknn /models/ +ADD https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/deci-fp16-yolonas_s-rk3566-v2.0.0-1.rknn /models/ +ADD https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/deci-fp16-yolonas_s-rk3568-v2.0.0-1.rknn /models/ +ADD https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/deci-fp16-yolonas_s-rk3576-v2.0.0-1.rknn /models/ +ADD https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/deci-fp16-yolonas_s-rk3588-v2.0.0-1.rknn /models/ RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffmpeg RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffprobe diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py index 614f51d84..40aaa1b85 100644 --- a/frigate/detectors/plugins/rknn.py +++ b/frigate/detectors/plugins/rknn.py @@ -15,9 +15,9 @@ DETECTOR_KEY = "rknn" supported_socs = ["rk3562", "rk3566", "rk3568", "rk3576", "rk3588"] -supported_models = ["^default-fp16-yolonas_[sml]$"] +supported_models = ["^deci-fp16-yolonas_[sml]$"] -default_model = "default-fp16-yolonas_s" +default_model = "deci-fp16-yolonas_s" model_chache_dir = "/config/model_cache/rknn_cache/" @@ -95,7 +95,7 @@ class Rknn(DetectionApi): model_props["filename"] = model_path + f"-{soc}-v2.0.0-1.rknn" if model_path == default_model: - model_path["path"] = "/models/" + model_props["filename"] + model_props["path"] = "/models/" + model_props["filename"] else: model_props["path"] = model_chache_dir + model_props["filename"] @@ -116,7 +116,7 @@ class Rknn(DetectionApi): os.mkdir(model_chache_dir) urllib.request.urlretrieve( - f"https://github.com/MarcA711/rknn-models/releases/tag/v2.0.0/{filename}", + f"https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/{filename}", model_chache_dir + filename, )