chris/frigate
1
0
Fork 0
forked from Github/frigate
Code Pull Requests Packages Activity

Reimplement support for rknn detector (#11365)

Browse Source 
* initial support for rknn detector

* remove purge_model_cache option

* update rknn

* support rk3576

* fix post_process_yolonas call

* add yolonas models

* update config

* exclude yolonas from image

* remove code
This commit is contained in:
Marc Altmann
2024-05-22 00:50:03 +02:00
committed by GitHub
parent 910c85b1c0
commit e91f3d8d9b
7 changed files with 231 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
frigate/detectors/detection_api.py
View File

@@ -72,7 +72,7 @@ class DetectionApi(ABC):
def post_process(self, output):
if self.detector_config.model.model_type == ModelTypeEnum.yolonas:
return self.yolonas(output)
return self.post_process_yolonas(output)
else:
raise ValueError(
f'Model type "{self.detector_config.model.model_type}" is currently not supported.'

185
frigate/detectors/plugins/rknn.py
View File

@@ -1,118 +1,157 @@
import logging
import os.path
import re
import urllib.request
from typing import Literal
try:
from hide_warnings import hide_warnings
except: # noqa: E722
def hide_warnings(func):
pass
from pydantic import Field
from frigate.detectors.detection_api import DetectionApi
from frigate.detectors.detector_config import BaseDetectorConfig
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn"
supported_socs = ["rk3562", "rk3566", "rk3568", "rk3588"]
supported_socs = ["rk3562", "rk3566", "rk3568", "rk3576", "rk3588"]
supported_models = {ModelTypeEnum.yolonas: "^deci-fp16-yolonas_[sml]$"}
model_chache_dir = "/config/model_cache/rknn_cache/"
class RknnDetectorConfig(BaseDetectorConfig):
type: Literal[DETECTOR_KEY]
core_mask: int = Field(default=0, ge=0, le=7, title="Core mask for NPU.")
num_cores: int = Field(default=0, ge=0, le=3, title="Number of NPU cores to use.")
purge_model_cache: bool = Field(default=True)
class Rknn(DetectionApi):
type_key = DETECTOR_KEY
def __init__(self, config: RknnDetectorConfig):
# find out SoC
try:
with open("/proc/device-tree/compatible") as file:
soc = file.read().split(",")[-1].strip("\x00")
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
)
)
if not os.path.isfile("/usr/lib/librknnrt.so"):
if "rk356" in soc:
os.rename("/usr/lib/librknnrt_rk356x.so", "/usr/lib/librknnrt.so")
elif "rk3588" in soc:
os.rename("/usr/lib/librknnrt_rk3588.so", "/usr/lib/librknnrt.so")
self.core_mask = config.core_mask
super().__init__(config)
self.height = config.model.height
self.width = config.model.width
core_mask = 2**config.num_cores - 1
soc = self.get_soc()
if True:
os.makedirs("/config/model_cache/rknn", exist_ok=True)
model_props = self.parse_model_input(config.model.path, soc)
if (config.model.width != 320) or (config.model.height != 320):
logger.error(
"Make sure to set the model width and height to 320 in your config.yml."
)
raise Exception(
"Make sure to set the model width and height to 320 in your config.yml."
)
if model_props["preset"]:
config.model.model_type = model_props["model_type"]
if config.model.input_pixel_format != "bgr":
logger.error(
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.'
)
raise Exception(
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.'
)
if config.model.input_tensor != "nhwc":
logger.error(
'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
)
raise Exception(
'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
)
if model_props["model_type"] == ModelTypeEnum.yolonas:
logger.info("""
You are using yolo-nas with weights from DeciAI.
These weights are subject to their license and can't be used commercially.
For more information, see: https://docs.deci.ai/super-gradients/latest/LICENSE.YOLONAS.html
""")
from rknnlite.api import RKNNLite
self.rknn = RKNNLite(verbose=False)
if self.rknn.load_rknn(self.model_path) != 0:
if self.rknn.load_rknn(model_props["path"]) != 0:
logger.error("Error initializing rknn model.")
if self.rknn.init_runtime(core_mask=self.core_mask) != 0:
if self.rknn.init_runtime(core_mask=core_mask) != 0:
logger.error(
"Error initializing rknn runtime. Do you run docker in privileged mode?"
)
raise Exception(
"RKNN does not currently support any models. Please see the docs for more info."
)
def __del__(self):
self.rknn.release()
@hide_warnings
def inference(self, tensor_input):
return self.rknn.inference(inputs=tensor_input)
def get_soc(self):
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.")
if soc not in supported_socs:
raise Exception(
f"Your SoC is not supported. Your SoC is: {soc}. Currently these SoCs are supported: {supported_socs}."
)
return soc
def parse_model_input(self, model_path, soc):
model_props = {}
# find out if user provides his own model
# user provided models should be a path and contain a "/"
if "/" in model_path:
model_props["preset"] = False
model_props["path"] = model_path
else:
model_props["preset"] = True
"""
Filenames follow this pattern:
origin-quant-basename-soc-tk_version-rev.rknn
origin: From where comes the model? default: upstream repo; rknn: modifications from airockchip
quant: i8 or fp16
basename: e.g. yolonas_s
soc: e.g. rk3588
tk_version: e.g. v2.0.0
rev: e.g. 1
Full name could be: default-fp16-yolonas_s-rk3588-v2.0.0-1.rknn
"""
model_matched = False
for model_type, pattern in supported_models.items():
if re.match(pattern, model_path):
model_matched = True
model_props["model_type"] = model_type
if model_matched:
model_props["filename"] = model_path + f"-{soc}-v2.0.0-1.rknn"
model_props["path"] = model_chache_dir + model_props["filename"]
if not os.path.isfile(model_props["path"]):
self.download_model(model_props["filename"])
else:
supported_models_str = ", ".join(
model[1:-1] for model in supported_models
)
raise Exception(
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)
urllib.request.urlretrieve(
f"https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/{filename}",
model_chache_dir + filename,
)
def check_config(self, config):
if (config.model.width != 320) or (config.model.height != 320):
raise Exception(
"Make sure to set the model width and height to 320 in your config.yml."
)
if config.model.input_pixel_format != "bgr":
raise Exception(
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.'
)
if config.model.input_tensor != "nhwc":
raise Exception(
'Make sure to set the model input_tensor to "nhwc" in your config.yml.'
)
def detect_raw(self, tensor_input):
output = self.inference(
output = self.rknn.inference(
[
tensor_input,
]
)
return self.postprocess(output[0])
return self.post_process(output)