Pull from event DB directly

frigate/data_processing/post/review_descriptions.py

@@ -34,7 +34,7 @@ from frigate.const import (
 from frigate.data_processing.types import PostProcessDataEnum
 from frigate.genai import GenAIClient
 from frigate.genai.manager import GenAIClientManager
-from frigate.models import Recordings, ReviewSegment
+from frigate.models import Event, Recordings, ReviewSegment
 from frigate.util.builtin import EventsPerSecond, InferenceSpeed
 from frigate.util.image import get_image_from_recording
 from frigate.util.recording_coverage import recording_overlap
@@ -439,13 +439,22 @@ class ReviewDescriptionProcessor(PostProcessorApi):
         if camera_config.review.genai.debug_save_thumbnails:
             self.save_debug_recording_frames(review_id, frames)
 
-        self.start_analysis(camera_config, final_data, frames)
+        self.start_analysis(
+            camera_config,
+            final_data,
+            frames,
+            get_named_objects(
+                final_data["data"].get("detections") or [],
+                self.config.all_attributes,
+            ),
+        )
 
     def start_analysis(
         self,
         camera_config: CameraConfig,
         final_data: dict[str, Any],
         frames: list[tuple[bytes, float]],
+        named_objects: list[tuple[str, str]] | None = None,
     ) -> None:
         """Kick off description generation for a review item in the background."""
         thumbs = [frame for frame, _ in frames]
@@ -481,6 +490,7 @@ class ReviewDescriptionProcessor(PostProcessorApi):
                 camera_config.review.genai,
                 sorted(self.config.all_labels),
                 self.config.all_attributes,
+                named_objects,
             ),
         ).start()
 
@@ -741,6 +751,29 @@ def format_classification_state_changes(
     return lines
 
 
+def get_named_objects(
+    detection_ids: list[str], attribute_labels: list[str]
+) -> list[tuple[str, str]]:
+    """Get the label and sub label of each named tracked object in a review item.
+
+    Editing a sub label after a review item ends rebuilds the stored sub_labels
+    without verified_objects, so the stored lists can't be paired by position.
+    """
+    events = (
+        Event.select(Event.label, Event.sub_label)
+        .where(Event.id.in_(detection_ids), Event.sub_label.is_null(False))
+        .order_by(Event.start_time)
+    )
+
+    # attribute sub labels (eg, delivery logos) are not names, which matches
+    # how the review maintainer records them
+    return [
+        (event.label, event.sub_label)
+        for event in events
+        if event.sub_label and event.sub_label not in attribute_labels
+    ]
+
+
 def run_analysis(
     requestor: InterProcessRequestor,
     genai_client: GenAIClient,
@@ -752,6 +785,7 @@ def run_analysis(
     genai_config: GenAIReviewConfig,
     labelmap_objects: list[str],
     attribute_labels: list[str],
+    named_objects: list[tuple[str, str]] | None,
 ) -> None:
     start = datetime.datetime.now().timestamp()
 
@@ -776,18 +810,20 @@ def run_analysis(
     unified_objects = []
 
     objects_list = final_data["data"]["objects"]
-    sub_labels_list = final_data["data"]["sub_labels"]
 
-    for i, verified_label in enumerate(final_data["data"]["verified_objects"]):
-        object_type = verified_label.replace("-verified", "").replace("_", " ")
+    if named_objects is None:
+        # live review items record a sub label for every verified object
+        named_objects = [
+            (verified_label.replace("-verified", ""), sub_label)
+            for verified_label, sub_label in zip(
+                final_data["data"]["verified_objects"],
+                final_data["data"]["sub_labels"],
+            )
+        ]
 
-        # editing a sub label after the review item ended rebuilds sub_labels
-        # without verified_objects, so a stored item can have fewer sub labels
-        if i >= len(sub_labels_list):
-            unified_objects.append(object_type)
-            continue
-
-        name = titlecase(sub_labels_list[i].replace("_", " "))
+    for label, sub_label in named_objects:
+        object_type = label.replace("_", " ")
+        name = titlecase(sub_label.replace("_", " "))
         unified_objects.append(f"{name} ← {object_type}")
 
     for label in objects_list:

frigate/test/test_review_sub_label_sync.py

@@ -3,8 +3,14 @@
 import unittest
 from unittest.mock import MagicMock
 
+from peewee import SqliteDatabase
+
 from frigate.config import FrigateConfig
-from frigate.data_processing.post.review_descriptions import run_analysis
+from frigate.data_processing.post.review_descriptions import (
+    get_named_objects,
+    run_analysis,
+)
+from frigate.models import Event
 
 CONFIG = """
 mqtt:
@@ -23,7 +29,9 @@ cameras:
 """
 
 
-def unified_objects(review_data: dict) -> list[str]:
+def unified_objects(
+    review_data: dict, named_objects: list[tuple[str, str]] | None = None
+) -> list[str]:
     """Run the review analysis and return the objects sent to the GenAI client."""
     camera_config = FrigateConfig.parse_yaml(CONFIG).cameras["front_door"]
     client = MagicMock()
@@ -38,7 +46,8 @@ def unified_objects(review_data: dict) -> list[str]:
         [],
         camera_config.review.genai,
         ["car", "person"],
-        [],
+        ["amazon"],
+        named_objects,
     )
     return client.generate_review_description.call_args.args[0]["unified_objects"]
 
@@ -57,20 +66,93 @@ class TestRunAnalysisSubLabels(unittest.TestCase):
 
         self.assertEqual(objects, ["Nick M ← person", "car"])
 
-    def test_missing_sub_label_falls_back_to_object_type(self):
-        # clearing a sub label on an ended review item leaves verified_objects
-        # stale, which previously raised an IndexError on regenerate
+
+class TestRegeneratedSubLabels(unittest.TestCase):
+    """Review items whose stored lists were rebuilt after a sub label edit."""
+
+    def setUp(self):
+        self.db = SqliteDatabase(":memory:")
+        self.db.bind([Event])
+        self.db.create_tables([Event])
+
+    def tearDown(self):
+        self.db.close()
+
+    def _insert_event(
+        self, id: str, label: str, sub_label: str | None, start_time: float
+    ) -> None:
+        Event.create(
+            id=id,
+            label=label,
+            sub_label=sub_label,
+            camera="front_door",
+            start_time=start_time,
+            end_time=start_time + 10,
+            top_score=0.9,
+            score=0.9,
+            false_positive=False,
+            zones=[],
+            thumbnail="",
+            has_clip=True,
+            has_snapshot=True,
+            region=[],
+            box=[],
+            area=0,
+            retain_indefinitely=False,
+            plus_id="",
+            model_hash="",
+            detector_type="cpu",
+            model_type="ssd",
+            data={},
+        )
+
+    def test_cleared_sub_label_does_not_move_to_another_object(self):
+        self._insert_event("car1", "car", None, 1)
+        self._insert_event("person1", "person", "Nick", 2)
+
+        named = get_named_objects(["car1", "person1"], ["amazon"])
         objects = unified_objects(
             {
                 "zones": [],
-                "objects": ["person-verified"],
+                "objects": ["car", "person-verified"],
+                "verified_objects": ["car-verified", "person-verified"],
+                "sub_labels": ["Nick"],
+                "classification_state_changes": [],
+            },
+            named,
+        )
+
+        self.assertEqual(objects, ["Nick ← person", "car"])
+
+    def test_repeated_name_is_kept_for_each_object(self):
+        # sub_labels are deduplicated when rebuilt, which previously raised an
+        # IndexError when the verified objects outnumbered them
+        self._insert_event("person1", "person", "Nick", 1)
+        self._insert_event("person2", "person", "Nick", 2)
+
+        named = get_named_objects(["person1", "person2"], ["amazon"])
+        objects = unified_objects(
+            {
+                "zones": [],
+                "objects": ["person-verified", "person-verified"],
                 "verified_objects": ["person-verified", "person-verified"],
                 "sub_labels": ["Nick"],
                 "classification_state_changes": [],
-            }
+            },
+            named,
         )
 
-        self.assertEqual(objects, ["Nick ← person", "person"])
+        self.assertEqual(objects, ["Nick ← person", "Nick ← person"])
+
+    def test_attribute_sub_labels_are_not_names(self):
+        self._insert_event("car1", "car", "amazon", 1)
+        self._insert_event("person1", "person", "Nick", 2)
+        self._insert_event("other", "person", "Bob", 3)
+
+        self.assertEqual(
+            get_named_objects(["car1", "person1"], ["amazon"]),
+            [("person", "Nick")],
+        )
 
 
 if __name__ == "__main__":