From 5f7740895624e58663e02e805932f815d1f08d5e Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Wed, 16 Oct 2024 15:22:34 -0600
Subject: [PATCH] Update logos handling (#14396)

* Add attribute for logos

* Clean up tracked object to pass model data

* Update default attributes map
---
 frigate/const.py                     | 16 +++++++++++++++-
 frigate/detectors/detector_config.py | 11 +++++++++++
 frigate/object_processing.py         | 17 +++++++----------
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/frigate/const.py b/frigate/const.py
index ad1aacd0f..c83b10e73 100644
--- a/frigate/const.py
+++ b/frigate/const.py
@@ -17,7 +17,21 @@ PLUS_API_HOST = "https://api.frigate.video"
 
 DEFAULT_ATTRIBUTE_LABEL_MAP = {
     "person": ["amazon", "face"],
-    "car": ["amazon", "fedex", "license_plate", "ups"],
+    "car": [
+        "amazon",
+        "an_post",
+        "dhl",
+        "dpd",
+        "fedex",
+        "gls",
+        "license_plate",
+        "nzpost",
+        "postnl",
+        "postnord",
+        "purolator",
+        "ups",
+        "usps",
+    ],
 }
 LABEL_CONSOLIDATION_MAP = {
     "car": 0.8,
diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py
index 90937d8f4..c40ef65bf 100644
--- a/frigate/detectors/detector_config.py
+++ b/frigate/detectors/detector_config.py
@@ -59,6 +59,7 @@ class ModelConfig(BaseModel):
     _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
     _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()
     _all_attributes: list[str] = PrivateAttr()
+    _all_attribute_logos: list[str] = PrivateAttr()
     _model_hash: str = PrivateAttr()
 
     @property
@@ -73,6 +74,10 @@ class ModelConfig(BaseModel):
     def all_attributes(self) -> list[str]:
         return self._all_attributes
 
+    @property
+    def all_attribute_logos(self) -> list[str]:
+        return self._all_attribute_logos
+
     @property
     def model_hash(self) -> str:
         return self._model_hash
@@ -93,6 +98,9 @@ class ModelConfig(BaseModel):
             unique_attributes.update(attributes)
 
         self._all_attributes = list(unique_attributes)
+        self._all_attribute_logos = list(
+            unique_attributes - set(["face", "license_plate"])
+        )
 
     def check_and_load_plus_model(
         self, plus_api: PlusApi, detector: str = None
@@ -140,6 +148,9 @@ class ModelConfig(BaseModel):
             unique_attributes.update(attributes)
 
         self._all_attributes = list(unique_attributes)
+        self._all_attribute_logos = list(
+            unique_attributes - set(["face", "license_plate"])
+        )
 
         self._merged_labelmap = {
             **{int(key): val for key, val in model_info["labelMap"].items()},
diff --git a/frigate/object_processing.py b/frigate/object_processing.py
index a5b06b76f..6e63562a4 100644
--- a/frigate/object_processing.py
+++ b/frigate/object_processing.py
@@ -20,6 +20,7 @@ from frigate.comms.inter_process import InterProcessRequestor
 from frigate.config import (
     CameraConfig,
     FrigateConfig,
+    ModelConfig,
     MqttConfig,
     RecordConfig,
     SnapshotsConfig,
@@ -109,8 +110,7 @@ def is_better_thumbnail(label, current_thumb, new_obj, frame_shape) -> bool:
 class TrackedObject:
     def __init__(
         self,
-        camera,
-        colormap,
+        model_config: ModelConfig,
         camera_config: CameraConfig,
         frame_cache,
         obj_data: dict[str, any],
@@ -120,8 +120,8 @@ class TrackedObject:
         del obj_data["score_history"]
 
         self.obj_data = obj_data
-        self.camera = camera
-        self.colormap = colormap
+        self.colormap = model_config.colormap
+        self.logos = model_config.all_attribute_logos
         self.camera_config = camera_config
         self.frame_cache = frame_cache
         self.zone_presence: dict[str, int] = {}
@@ -245,9 +245,7 @@ class TrackedObject:
         # populate the sub_label for object with highest scoring logo
         if self.obj_data["label"] in ["car", "package", "person"]:
             recognized_logos = {
-                k: self.attributes[k]
-                for k in ["ups", "fedex", "amazon"]
-                if k in self.attributes
+                k: self.attributes[k] for k in self.logos if k in self.attributes
             }
             if len(recognized_logos) > 0:
                 max_logo = max(recognized_logos, key=recognized_logos.get)
@@ -291,7 +289,7 @@ class TrackedObject:
     def to_dict(self, include_thumbnail: bool = False):
         event = {
             "id": self.obj_data["id"],
-            "camera": self.camera,
+            "camera": self.camera_config.name,
             "frame_time": self.obj_data["frame_time"],
             "snapshot": self.thumbnail_data,
             "label": self.obj_data["label"],
@@ -707,8 +705,7 @@ class CameraState:
 
         for id in new_ids:
             new_obj = tracked_objects[id] = TrackedObject(
-                self.name,
-                self.config.model.colormap,
+                self.config.model,
                 self.camera_config,
                 self.frame_cache,
                 current_detections[id],