From de3415f08d52d90872aec42b3be4ed4014d86fac Mon Sep 17 00:00:00 2001 From: eugene_yao Date: Mon, 17 Jun 2024 14:27:21 -0400 Subject: [PATCH] "Update parse_model_input function to include config information; add a condition to check for specified model_type when using custom model path." --- frigate/detectors/plugins/rknn.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/frigate/detectors/plugins/rknn.py b/frigate/detectors/plugins/rknn.py index 49f1f4c02..c1a2977b3 100644 --- a/frigate/detectors/plugins/rknn.py +++ b/frigate/detectors/plugins/rknn.py @@ -36,7 +36,7 @@ class Rknn(DetectionApi): core_mask = 2**config.num_cores - 1 soc = self.get_soc() - model_props = self.parse_model_input(config.model.path, soc) + model_props = self.parse_model_input(config, soc) if model_props["preset"]: config.model.model_type = model_props["model_type"] @@ -75,7 +75,9 @@ class Rknn(DetectionApi): return soc - def parse_model_input(self, model_path, soc): + def parse_model_input(self, config, soc): + model_path = config.model.path + model_props = {} # find out if user provides his own model @@ -83,6 +85,12 @@ class Rknn(DetectionApi): if "/" in model_path: model_props["preset"] = False model_props["path"] = model_path + if config.model.model_type: + model_props["model_type"] = config.model.model_type + else: + raise Exception( + "You must specify model_type if specifying your own model file." + ) else: model_props["preset"] = True