From 1c85f774eb87e5df76eedf194361ff36bfe31ea9 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Mon, 16 Aug 2021 08:02:04 -0500
Subject: [PATCH] move colormap to config

---
 frigate/config.py            | 17 ++++++++++++++---
 frigate/object_processing.py | 33 +++++++++++++++------------------
 frigate/process_clip.py      |  2 +-
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/frigate/config.py b/frigate/config.py
index e5fdff5d9..ea6ea3280 100644
--- a/frigate/config.py
+++ b/frigate/config.py
@@ -1,18 +1,18 @@
 from __future__ import annotations
 
-from enum import Enum
 import json
 import logging
 import os
+from enum import Enum
 from typing import Dict, List, Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
+import yaml
 from pydantic import BaseModel, Field, validator
 from pydantic.fields import PrivateAttr
-import yaml
 
-from frigate.const import BASE_DIR, RECORD_DIR, CACHE_DIR
+from frigate.const import BASE_DIR, CACHE_DIR, RECORD_DIR
 from frigate.edgetpu import load_labels
 from frigate.util import create_mask, deep_merge
 
@@ -578,11 +578,16 @@ class ModelConfig(BaseModel):
         default_factory=dict, title="Labelmap customization."
     )
     _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
+    _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()
 
     @property
     def merged_labelmap(self) -> Dict[int, str]:
         return self._merged_labelmap
 
+    @property
+    def colormap(self) -> Dict[int, tuple[int, int, int]]:
+        return self._colormap
+
     def __init__(self, **config):
         super().__init__(**config)
 
@@ -591,6 +596,12 @@ class ModelConfig(BaseModel):
             **config.get("labelmap", {}),
         }
 
+        cmap = plt.cm.get_cmap("tab10", len(self._merged_labelmap.keys()))
+
+        self._colormap = {}
+        for key, val in self._merged_labelmap.items():
+            self._colormap[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3])
+
 
 class LogLevelEnum(str, Enum):
     debug = "debug"
diff --git a/frigate/object_processing.py b/frigate/object_processing.py
index 999b67e6e..4ebd20870 100644
--- a/frigate/object_processing.py
+++ b/frigate/object_processing.py
@@ -1,5 +1,5 @@
-import copy
 import base64
+import copy
 import datetime
 import hashlib
 import itertools
@@ -14,30 +14,20 @@ from statistics import mean, median
 from typing import Callable, Dict
 
 import cv2
-import matplotlib.pyplot as plt
 import numpy as np
 
-from frigate.config import FrigateConfig, CameraConfig
-from frigate.const import RECORD_DIR, CLIPS_DIR, CACHE_DIR
+from frigate.config import CameraConfig, FrigateConfig
+from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
 from frigate.edgetpu import load_labels
 from frigate.util import (
     SharedMemoryFrameManager,
+    calculate_region,
     draw_box_with_label,
     draw_timestamp,
-    calculate_region,
 )
 
 logger = logging.getLogger(__name__)
 
-PATH_TO_LABELS = "/labelmap.txt"
-
-LABELS = load_labels(PATH_TO_LABELS)
-cmap = plt.cm.get_cmap("tab10", len(LABELS.keys()))
-
-COLOR_MAP = {}
-for key, val in LABELS.items():
-    COLOR_MAP[val] = tuple(int(round(255 * c)) for c in cmap(key)[:3])
-
 
 def on_edge(box, frame_shape):
     if (
@@ -72,9 +62,12 @@ def is_better_thumbnail(current_thumb, new_obj, frame_shape) -> bool:
 
 
 class TrackedObject:
-    def __init__(self, camera, camera_config: CameraConfig, frame_cache, obj_data):
+    def __init__(
+        self, camera, colormap, camera_config: CameraConfig, frame_cache, obj_data
+    ):
         self.obj_data = obj_data
         self.camera = camera
+        self.colormap = colormap
         self.camera_config = camera_config
         self.frame_cache = frame_cache
         self.current_zones = []
@@ -247,7 +240,7 @@ class TrackedObject:
 
         if bounding_box:
             thickness = 2
-            color = COLOR_MAP[self.obj_data["label"]]
+            color = self.colormap[self.obj_data["label"]]
 
             # draw the bounding boxes on the frame
             box = self.thumbnail_data["box"]
@@ -357,7 +350,7 @@ class CameraState:
             for obj in tracked_objects.values():
                 if obj["frame_time"] == frame_time:
                     thickness = 2
-                    color = COLOR_MAP[obj["label"]]
+                    color = self.config.model.colormap[obj["label"]]
                 else:
                     thickness = 1
                     color = (255, 0, 0)
@@ -448,7 +441,11 @@ class CameraState:
 
         for id in new_ids:
             new_obj = tracked_objects[id] = TrackedObject(
-                self.name, self.camera_config, self.frame_cache, current_detections[id]
+                self.name,
+                self.config.model.colormap,
+                self.camera_config,
+                self.frame_cache,
+                current_detections[id],
             )
 
             # call event handlers
diff --git a/frigate/process_clip.py b/frigate/process_clip.py
index b6462121c..ee9240338 100644
--- a/frigate/process_clip.py
+++ b/frigate/process_clip.py
@@ -14,7 +14,7 @@ import numpy as np
 from frigate.config import FRIGATE_CONFIG_SCHEMA, FrigateConfig
 from frigate.edgetpu import LocalObjectDetector
 from frigate.motion import MotionDetector
-from frigate.object_processing import COLOR_MAP, CameraState
+from frigate.object_processing import CameraState
 from frigate.objects import ObjectTracker
 from frigate.util import (
     DictFrameManager,