From 2185456ea5c57ac27fdd6fc82a7edcc0a32f5918 Mon Sep 17 00:00:00 2001 From: Nick Mowen Date: Fri, 6 Jan 2023 12:29:40 -0700 Subject: [PATCH] Only create colormap for enabled labels --- frigate/config.py | 7 +++++++ frigate/detectors/detector_config.py | 10 ++++++---- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/frigate/config.py b/frigate/config.py index 66712406f..54d5f3fc4 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -962,6 +962,12 @@ class FrigateConfig(FrigateBaseModel): camera_config.create_ffmpeg_cmds() config.cameras[name] = camera_config + # get list of unique enabled labels for tracking + enabled_labels = set(config.objects.track) + + for _, camera in config.cameras.items(): + enabled_labels.update(camera.objects.track) + for key, detector in config.detectors.items(): detector_config: DetectorConfig = parse_obj_as(DetectorConfig, detector) if detector_config.model is None: @@ -986,6 +992,7 @@ class FrigateConfig(FrigateBaseModel): config.model.dict(exclude_unset=True), ) detector_config.model = ModelConfig.parse_obj(merged_model) + detector_config.model.create_colormap(enabled_labels) config.detectors[key] = detector_config return config diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index 7eb8701f3..747a12de4 100644 --- a/frigate/detectors/detector_config.py +++ b/frigate/detectors/detector_config.py @@ -55,11 +55,13 @@ class ModelConfig(BaseModel): **load_labels(config.get("labelmap_path", "/labelmap.txt")), **config.get("labelmap", {}), } - - cmap = plt.cm.get_cmap("tab10", len(self._merged_labelmap.keys())) - self._colormap = {} - for key, val in self._merged_labelmap.items(): + + def create_colormap(self, enabled_labels: set[str]) -> None: + """Get a list of colors for enabled labels.""" + cmap = plt.cm.get_cmap("tab10", len(enabled_labels)) + + for key, val in enumerate(enabled_labels): self._colormap[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3]) class Config: