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,11 +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"]):
+    if named_objects is None:
-        object_type = verified_label.replace("-verified", "").replace("_", " ")
+        # live review items record a sub label for every verified object
-        name = titlecase(sub_labels_list[i].replace("_", " "))
+        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"],
+            )
+        ]
+
+    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

@@ -0,0 +1,159 @@
+"""Tests for pairing review sub labels with their verified objects."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from peewee import SqliteDatabase
+
+from frigate.config import FrigateConfig
+from frigate.data_processing.post.review_descriptions import (
+    get_named_objects,
+    run_analysis,
+)
+from frigate.models import Event
+
+CONFIG = """
+mqtt:
+  enabled: False
+cameras:
+  front_door:
+    ffmpeg:
+      inputs:
+        - path: rtsp://10.0.0.1:554/video
+          roles:
+            - detect
+    detect:
+      width: 1920
+      height: 1080
+      fps: 5
+"""
+
+
+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()
+    client.generate_review_description.return_value = None
+    run_analysis(
+        MagicMock(),
+        client,
+        MagicMock(),
+        camera_config,
+        {"id": "review1", "start_time": 0, "end_time": 10, "data": review_data},
+        [],
+        [],
+        camera_config.review.genai,
+        ["car", "person"],
+        ["amazon"],
+        named_objects,
+    )
+    return client.generate_review_description.call_args.args[0]["unified_objects"]
+
+
+class TestRunAnalysisSubLabels(unittest.TestCase):
+    def test_verified_objects_are_named_by_their_sub_label(self):
+        objects = unified_objects(
+            {
+                "zones": [],
+                "objects": ["person-verified", "car"],
+                "verified_objects": ["person-verified"],
+                "sub_labels": ["nick_m"],
+                "classification_state_changes": [],
+            }
+        )
+
+        self.assertEqual(objects, ["Nick M ← person", "car"])
+
+
+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": ["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", "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__":
+    unittest.main()