From 543cad54976b9cc4909239a61713b283c9ea8435 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Fri, 6 Jan 2023 18:31:54 -0700 Subject: [PATCH] Only set colors for enabled objects (#4936) * Only create colormap for enabled labels * Fix assigning --- frigate/config.py | 8 ++++++++ frigate/detectors/detector_config.py | 10 ++++++---- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/frigate/config.py b/frigate/config.py index 66712406f..e55db040e 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -962,6 +962,14 @@ 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) + + config.model.create_colormap(enabled_labels) + for key, detector in config.detectors.items(): detector_config: DetectorConfig = parse_obj_as(DetectorConfig, detector) if detector_config.model is None: 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: