-
-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Reimplement support for rknn detector (#11365)
* 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
- Loading branch information
Showing
7 changed files
with
228 additions
and
87 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 | ||
super().__init__(config) | ||
self.height = config.model.height | ||
self.width = config.model.width | ||
core_mask = 2**config.num_cores - 1 | ||
soc = self.get_soc() | ||
|
||
model_props = self.parse_model_input(config.model.path, soc) | ||
|
||
if model_props["preset"]: | ||
config.model.model_type = model_props["model_type"] | ||
|
||
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(model_props["path"]) != 0: | ||
logger.error("Error initializing rknn model.") | ||
if self.rknn.init_runtime(core_mask=core_mask) != 0: | ||
logger.error( | ||
"Error initializing rknn runtime. Do you run docker in privileged mode?" | ||
) | ||
|
||
def __del__(self): | ||
self.rknn.release() | ||
|
||
def get_soc(self): | ||
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 | ||
) | ||
f"Your SoC is not supported. Your SoC is: {soc}. Currently these SoCs are supported: {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 | ||
self.height = config.model.height | ||
self.width = config.model.width | ||
|
||
if True: | ||
os.makedirs("/config/model_cache/rknn", exist_ok=True) | ||
|
||
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." | ||
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( | ||
"Make sure to set the model width and height to 320 in your config.yml." | ||
f"Model {model_path} is unsupported. Provide your own model or choose one of the following: {supported_models_str}" | ||
) | ||
|
||
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.' | ||
) | ||
return model_props | ||
|
||
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.' | ||
) | ||
def download_model(self, filename): | ||
if not os.path.isdir(model_chache_dir): | ||
os.mkdir(model_chache_dir) | ||
|
||
from rknnlite.api import RKNNLite | ||
urllib.request.urlretrieve( | ||
f"https://github.com/MarcA711/rknn-models/releases/download/v2.0.0/{filename}", | ||
model_chache_dir + filename, | ||
) | ||
|
||
self.rknn = RKNNLite(verbose=False) | ||
if self.rknn.load_rknn(self.model_path) != 0: | ||
logger.error("Error initializing rknn model.") | ||
if self.rknn.init_runtime(core_mask=self.core_mask) != 0: | ||
logger.error( | ||
"Error initializing rknn runtime. Do you run docker in privileged mode?" | ||
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." | ||
) | ||
|
||
raise Exception( | ||
"RKNN does not currently support any models. Please see the docs for more info." | ||
) | ||
|
||
def __del__(self): | ||
self.rknn.release() | ||
if config.model.input_pixel_format != "bgr": | ||
raise Exception( | ||
'Make sure to set the model input_pixel_format to "bgr" in your config.yml.' | ||
) | ||
|
||
@hide_warnings | ||
def inference(self, tensor_input): | ||
return self.rknn.inference(inputs=tensor_input) | ||
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) |