From b466a951e3c0895d6b9b22378f1941552f577160 Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Fri, 16 Jun 2023 15:32:29 -0500
Subject: [PATCH] pass attribute labels as attributes

---
 frigate/object_processing.py | 16 ++++++++++++++
 frigate/video.py             | 43 ++++++++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/frigate/object_processing.py b/frigate/object_processing.py
index 6455f6eba..923e111ab 100644
--- a/frigate/object_processing.py
+++ b/frigate/object_processing.py
@@ -214,6 +214,7 @@ class TrackedObject:
             "entered_zones": self.entered_zones.copy(),
             "has_clip": self.has_clip,
             "has_snapshot": self.has_snapshot,
+            "attributes": self.obj_data["attributes"],
         }
 
         if include_thumbnail:
@@ -421,6 +422,21 @@ class CameraState:
                     color=color,
                 )
 
+                # draw any attributes
+                for attribute in obj["attributes"]:
+                    box = attribute["box"]
+                    draw_box_with_label(
+                        frame_copy,
+                        box[0],
+                        box[1],
+                        box[2],
+                        box[3],
+                        attribute["label"],
+                        f"{attribute['score']:.0%}",
+                        thickness=thickness,
+                        color=color,
+                    )
+
         if draw_options.get("regions"):
             for region in regions:
                 cv2.rectangle(
diff --git a/frigate/video.py b/frigate/video.py
index 25ae35664..c02ad15c4 100755
--- a/frigate/video.py
+++ b/frigate/video.py
@@ -723,6 +723,14 @@ def process_frames(
     stop_event,
     exit_on_empty: bool = False,
 ):
+    # attribute labels are not tracked and are not assigned regions
+    attribute_label_map = {
+        "person": ["face", "amazon"],
+        "car": ["ups", "fedex", "amazon", "license_plate"],
+    }
+    all_attribute_labels = [
+        item for sublist in attribute_label_map.values() for item in sublist
+    ]
     fps = process_info["process_fps"]
     detection_fps = process_info["detection_fps"]
     current_frame_time = process_info["detection_frame"]
@@ -758,6 +766,7 @@ def process_frames(
         motion_boxes = motion_detector.detect(frame) if motion_enabled.value else []
 
         regions = []
+        consolidated_detections = []
 
         # if detection is disabled
         if not detection_enabled.value:
@@ -894,12 +903,42 @@ def process_frames(
                 consolidated_detections = get_consolidated_object_detections(
                     detected_object_groups
                 )
+                tracked_detections = [
+                    d
+                    for d in consolidated_detections
+                    if d[0] not in all_attribute_labels
+                ]
                 # now that we have refined our detections, we need to track objects
-                object_tracker.match_and_update(frame_time, consolidated_detections)
+                object_tracker.match_and_update(frame_time, tracked_detections)
             # else, just update the frame times for the stationary objects
             else:
                 object_tracker.update_frame_times(frame_time)
 
+        # group the attribute detections based on what label they apply to
+        attribute_detections = {}
+        for label, attribute_labels in attribute_label_map.items():
+            attribute_detections[label] = [
+                d for d in consolidated_detections if d[0] in attribute_labels
+            ]
+
+        # build detections and add attributes
+        detections = {}
+        for obj in object_tracker.tracked_objects.values():
+            attributes = []
+            # if the objects label has associated attribute detections
+            if obj["label"] in attribute_detections.keys():
+                # add them to attributes if they intersect
+                for attribute_detection in attribute_detections[obj["label"]]:
+                    if box_inside(obj["box"], (attribute_detection[2])):
+                        attributes.append(
+                            {
+                                "label": attribute_detection[0],
+                                "score": attribute_detection[1],
+                                "box": attribute_detection[2],
+                            }
+                        )
+            detections[obj["id"]] = {**obj, "attributes": attributes}
+
         # debug object tracking
         if False:
             bgr_frame = cv2.cvtColor(
@@ -982,7 +1021,7 @@ def process_frames(
                 (
                     camera_name,
                     frame_time,
-                    object_tracker.tracked_objects,
+                    detections,
                     motion_boxes,
                     regions,
                 )