diff --git a/frigate/data_processing/post/review_descriptions.py b/frigate/data_processing/post/review_descriptions.py index 431472e6f4..d2f5619e12 100644 --- a/frigate/data_processing/post/review_descriptions.py +++ b/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"]): - object_type = verified_label.replace("-verified", "").replace("_", " ") - name = titlecase(sub_labels_list[i].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"], + ) + ] + + 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: diff --git a/frigate/test/test_review_sub_label_sync.py b/frigate/test/test_review_sub_label_sync.py new file mode 100644 index 0000000000..9a22a8a5f1 --- /dev/null +++ b/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()