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

Include libraries and .rknn models for other Rockchip SoCs (#8649)

Browse Source 
* support for other yolov models and config checks

* apply code formatting

* Information about core mask and inference speed

* update rknn postprocess and remove params

* update model selection

* Apply suggestions from code review

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

* support rknn on all socs

* apply changes from review and fix post process bug

* apply code formatting

* update tip in object_detectors docs

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
This commit is contained in:
Marc Altmann
2023-11-18 14:53:49 +01:00
committed by GitHub
parent 2da99c2308
commit c6208b266b
5 changed files with 72 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -22,7 +22,9 @@ logger = logging.getLogger(__name__)
DETECTOR_KEY = "rknn"
yolov8_rknn_models = {
supported_socs = ["rk3562", "rk3566", "rk3568", "rk3588"]
yolov8_suffix = {
"default-yolov8n": "n",
"default-yolov8s": "s",
"default-yolov8m": "m",
@@ -40,28 +42,57 @@ 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 "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.model_path = config.model.path or "default-yolov8n"
self.core_mask = config.core_mask
self.height = config.model.height
self.width = config.model.width
if self.model_path in yolov8_rknn_models:
if self.model_path in yolov8_suffix:
if self.model_path == "default-yolov8n":
self.model_path = "/models/yolov8n-320x320.rknn"
self.model_path = "/models/rknn/yolov8n-320x320-{soc}.rknn".format(
soc=soc
)
else:
model_suffix = yolov8_rknn_models[self.model_path]
model_suffix = yolov8_suffix[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,
)
@@ -140,10 +171,10 @@ class Rknn(DetectionApi):
boxes = np.transpose(
np.vstack(
(
results[:, 1] - 0.5 * results[:, 3],
results[:, 0] - 0.5 * results[:, 2],
results[:, 3] + 0.5 * results[:, 3],
results[:, 2] + 0.5 * results[:, 2],
(results[:, 1] - 0.5 * results[:, 3]) / self.height,
(results[:, 0] - 0.5 * results[:, 2]) / self.width,
(results[:, 1] + 0.5 * results[:, 3]) / self.height,
(results[:, 0] + 0.5 * results[:, 2]) / self.width,
)
)
)