diff --git a/frigate/data_processing/real_time/custom_classification.py b/frigate/data_processing/real_time/custom_classification.py index 19a40fd99..5a58cf122 100644 --- a/frigate/data_processing/real_time/custom_classification.py +++ b/frigate/data_processing/real_time/custom_classification.py @@ -97,7 +97,7 @@ class CustomStateClassificationProcessor(RealTimeProcessorApi): self.interpreter.allocate_tensors() self.tensor_input_details = self.interpreter.get_input_details() self.tensor_output_details = self.interpreter.get_output_details() - self.labelmap = load_labels(labelmap_path, prefill=0) + self.labelmap = load_labels(labelmap_path, prefill=0, indexed=False) self.classifications_per_second.start() def __update_metrics(self, duration: float) -> None: @@ -398,7 +398,7 @@ class CustomObjectClassificationProcessor(RealTimeProcessorApi): self.interpreter.allocate_tensors() self.tensor_input_details = self.interpreter.get_input_details() self.tensor_output_details = self.interpreter.get_output_details() - self.labelmap = load_labels(labelmap_path, prefill=0) + self.labelmap = load_labels(labelmap_path, prefill=0, indexed=False) def __update_metrics(self, duration: float) -> None: self.classifications_per_second.update() diff --git a/frigate/util/builtin.py b/frigate/util/builtin.py index b1a76214b..4443888cf 100644 --- a/frigate/util/builtin.py +++ b/frigate/util/builtin.py @@ -129,7 +129,7 @@ def get_ffmpeg_arg_list(arg: Any) -> list: return arg if isinstance(arg, list) else shlex.split(arg) -def load_labels(path: Optional[str], encoding="utf-8", prefill=91): +def load_labels(path: Optional[str], encoding="utf-8", prefill=91, indexed: bool | None = None): """Loads labels from file (with or without index numbers). Args: path: path to label file. @@ -146,11 +146,12 @@ def load_labels(path: Optional[str], encoding="utf-8", prefill=91): if not lines: return {} - if lines[0].split(" ", maxsplit=1)[0].isdigit(): + if indexed != False and lines[0].split(" ", maxsplit=1)[0].isdigit(): pairs = [line.split(" ", maxsplit=1) for line in lines] labels.update({int(index): label.strip() for index, label in pairs}) else: labels.update({index: line.strip() for index, line in enumerate(lines)}) + return labels