support rknn on all socs

This commit is contained in:
MarcA711 2023-11-17 21:11:40 +00:00
parent b7270a8b04
commit e0be59b3be
4 changed files with 52 additions and 13 deletions

View File

@ -10,10 +10,17 @@ RUN sed -i "/https:\/\//d" /requirements-wheels.txt
RUN pip3 wheel --wheel-dir=/rk-wheels -c /requirements-wheels.txt -r /requirements-wheels-rk.txt
FROM wget as rk-downloads
RUN wget -qO librknnrt.so https://github.com/MarcA711/rknpu2/raw/master/runtime/RK3588/Linux/librknn_api/aarch64/librknnrt.so
RUN wget -qO ffmpeg https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/latest/ffmpeg
RUN wget -qO ffprobe https://github.com/MarcA711/Rockchip-FFmpeg-Builds/releases/download/latest/ffprobe
RUN wget -qO yolov8n-320x320.rknn https://github.com/MarcA711/rknn-models/releases/download/latest/yolov8n-320x320.rknn
RUN wget -qO librknnrt_rk356x.so https://github.com/MarcA711/rknpu2/releases/download/v1.5.2/librknnrt_rk356x.so
RUN wget -qO librknnrt_rk3588.so https://github.com/MarcA711/rknpu2/releases/download/v1.5.2/librknnrt_rk3588.so
RUN wget -qO yolov8n-320x320-rk3562.rknn https://github.com/MarcA711/rknn-models/releases/download/v1.5.2-rk3562/yolov8n-320x320-rk3562.rknn
RUN wget -qO yolov8n-320x320-rk3566.rknn https://github.com/MarcA711/rknn-models/releases/download/v1.5.2-rk3566/yolov8n-320x320-rk3566.rknn
RUN wget -qO yolov8n-320x320-rk3568.rknn https://github.com/MarcA711/rknn-models/releases/download/v1.5.2-rk3568/yolov8n-320x320-rk3568.rknn
RUN wget -qO yolov8n-320x320-rk3588.rknn https://github.com/MarcA711/rknn-models/releases/download/v1.5.2-rk3588/yolov8n-320x320-rk3588.rknn
FROM deps AS rk-deps
ARG TARGETARCH
@ -22,8 +29,14 @@ RUN --mount=type=bind,from=rk-wheels,source=/rk-wheels,target=/deps/rk-wheels \
WORKDIR /opt/frigate/
COPY --from=rootfs / /
COPY --from=rk-downloads /rootfs/librknnrt.so /usr/lib/
COPY --from=rk-downloads /rootfs/yolov8n-320x320.rknn /models/
COPY --from=rk-downloads /rootfs/librknnrt_rk356x.so /usr/lib/
COPY --from=rk-downloads /rootfs/librknnrt_rk3588.so /usr/lib/
COPY --from=rk-downloads /rootfs/yolov8n-320x320-rk3562.rknn /models/
COPY --from=rk-downloads /rootfs/yolov8n-320x320-rk3566.rknn /models/
COPY --from=rk-downloads /rootfs/yolov8n-320x320-rk3568.rknn /models/
COPY --from=rk-downloads /rootfs/yolov8n-320x320-rk3588.rknn /models/
RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffmpeg
RUN rm -rf /usr/lib/btbn-ffmpeg/bin/ffprobe

View File

@ -1,2 +1,2 @@
hide-warnings == 0.17
rknn-toolkit-lite2 @ https://github.com/MarcA711/rknn-toolkit2/raw/master/rknn_toolkit_lite2/packages/rknn_toolkit_lite2-1.5.2-cp39-cp39-linux_aarch64.whl
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

View File

@ -295,9 +295,9 @@ To verify that the integration is working correctly, start Frigate and observe t
## Rockchip RKNN-Toolkit-Lite2
This detector is only available if one of the following Rockchip SoCs is used:
- RK3566/RK3568
- RK3588/RK3588S
- RV1103/RV1106
- RK3568
- RK3566
- RK3562
These SoCs come with a NPU that will highly speed up detection.

View File

@ -22,6 +22,8 @@ logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn"
supported_socs = ["rk3562", "rk3566", "rk3568", "rk3588"]
yolov8_rknn_models = {
"default-yolov8n": "n",
"default-yolov8s": "s",
@ -40,6 +42,28 @@ class Rknn(DetectionApi):
type_key = DETECTOR_KEY
def __init__(self, config: RknnDetectorConfig):
# find out SoC
try:
with open("/proc/device-tree/compatible") as file:
device_string = file.read()
device_string_parts = device_string.split(",")
soc = device_string_parts[-1]
except FileNotFoundError:
logger.error("Make sure to run docker in privileged mode.")
raise Exception("Make sure to run docker in privileged mode.")
if soc not in supported_socs:
logger.error(
"Your SoC is not supported. Your SoC is: {}. Currently these SoCs are supported: {}.".format(
soc, supported_socs
)
)
raise Exception(
"Your SoC is not supported. Your SoC is: {}. Currently these SoCs are supported: {}.".format(
soc, supported_socs
)
)
self.model_path = config.model.path or "default-yolov8n"
self.core_mask = config.core_mask
self.height = config.model.height
@ -47,21 +71,23 @@ class Rknn(DetectionApi):
if self.model_path in yolov8_rknn_models:
if self.model_path == "default-yolov8n":
self.model_path = "/models/yolov8n-320x320.rknn"
self.model_path = "/models/yolov8n-320x320-{soc}.rknn".format(soc=soc)
else:
model_suffix = yolov8_rknn_models[self.model_path]
self.model_path = (
"/config/model_cache/rknn/yolov8{}-320x320.rknn".format(
model_suffix
"/config/model_cache/rknn/yolov8{suffix}-320x320-{soc}.rknn".format(
suffix=model_suffix, soc=soc
)
)
os.makedirs("/config/model_cache/rknn", exist_ok=True)
if not os.path.isfile(self.model_path):
logger.info("Downloading yolov8{} model.".format(model_suffix))
logger.info(
"Downloading yolov8{suffix} model.".format(suffix=model_suffix)
)
urllib.request.urlretrieve(
"https://github.com/MarcA711/rknn-models/releases/download/latest/yolov8{}-320x320.rknn".format(
model_suffix
"https://github.com/MarcA711/rknn-models/releases/download/v1.5.2-{soc}/yolov8{suffix}-320x320-{soc}.rknn".format(
soc=soc, suffix=model_suffix
),
self.model_path,
)