diff --git a/frigate/api/defs/regenerate_query_parameters.py b/frigate/api/defs/regenerate_query_parameters.py new file mode 100644 index 000000000..4a8326cda --- /dev/null +++ b/frigate/api/defs/regenerate_query_parameters.py @@ -0,0 +1,7 @@ +from typing import Optional + +from pydantic import BaseModel + + +class RegenerateQueryParameters(BaseModel): + source: Optional[str] = "thumbnails" diff --git a/frigate/api/event.py b/frigate/api/event.py index 4e45b10de..3c861f901 100644 --- a/frigate/api/event.py +++ b/frigate/api/event.py @@ -31,6 +31,9 @@ from frigate.api.defs.events_query_parameters import ( EventsSearchQueryParams, EventsSummaryQueryParams, ) +from frigate.api.defs.regenerate_query_parameters import ( + RegenerateQueryParameters, +) from frigate.api.defs.tags import Tags from frigate.const import ( CLIPS_DIR, @@ -996,7 +999,9 @@ def set_description( @router.put("/events/{event_id}/description/regenerate") -def regenerate_description(request: Request, event_id: str): +def regenerate_description( + request: Request, event_id: str, params: RegenerateQueryParameters = Depends() +): try: event: Event = Event.get(Event.id == event_id) except DoesNotExist: @@ -1009,7 +1014,7 @@ def regenerate_description(request: Request, event_id: str): request.app.frigate_config.semantic_search.enabled and request.app.frigate_config.genai.enabled ): - request.app.event_metadata_updater.publish(event.id) + request.app.event_metadata_updater.publish((event.id, params.source)) return JSONResponse( content=( @@ -1017,7 +1022,8 @@ def regenerate_description(request: Request, event_id: str): "success": True, "message": "Event " + event_id - + " description regeneration has been requested.", + + " description regeneration has been requested using " + + params.source, } ), status_code=200, diff --git a/frigate/comms/event_metadata_updater.py b/frigate/comms/event_metadata_updater.py index d435b149e..aeede6d8e 100644 --- a/frigate/comms/event_metadata_updater.py +++ b/frigate/comms/event_metadata_updater.py @@ -4,6 +4,8 @@ import logging from enum import Enum from typing import Optional +from frigate.events.types import RegenerateDescriptionEnum + from .zmq_proxy import Publisher, Subscriber logger = logging.getLogger(__name__) @@ -23,6 +25,9 @@ class EventMetadataPublisher(Publisher): topic = topic.value super().__init__(topic) + def publish(self, payload: tuple[str, RegenerateDescriptionEnum]) -> None: + super().publish(payload) + class EventMetadataSubscriber(Subscriber): """Simplifies receiving event metadata.""" @@ -35,10 +40,12 @@ class EventMetadataSubscriber(Subscriber): def check_for_update( self, timeout: float = None - ) -> Optional[tuple[EventMetadataTypeEnum, any]]: + ) -> Optional[tuple[EventMetadataTypeEnum, str, RegenerateDescriptionEnum]]: return super().check_for_update(timeout) def _return_object(self, topic: str, payload: any) -> any: if payload is None: - return (None, None) - return (EventMetadataTypeEnum[topic[len(self.topic_base) :]], payload) + return (None, None, None) + topic = EventMetadataTypeEnum[topic[len(self.topic_base) :]] + event_id, source = payload + return (topic, event_id, RegenerateDescriptionEnum(source)) diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py index 31f5bbd27..c6add4eb3 100644 --- a/frigate/embeddings/maintainer.py +++ b/frigate/embeddings/maintainer.py @@ -172,13 +172,15 @@ class EmbeddingMaintainer(threading.Thread): def _process_event_metadata(self): # Check for regenerate description requests - (topic, event_id) = self.event_metadata_subscriber.check_for_update(timeout=1) + (topic, event_id, source) = self.event_metadata_subscriber.check_for_update( + timeout=1 + ) if topic is None: return if event_id: - self.handle_regenerate_description(event_id) + self.handle_regenerate_description(event_id, source) def _create_thumbnail(self, yuv_frame, box, height=500) -> Optional[bytes]: """Return jpg thumbnail of a region of the frame.""" @@ -241,7 +243,7 @@ class EmbeddingMaintainer(threading.Thread): description, ) - def handle_regenerate_description(self, event_id: str) -> None: + def handle_regenerate_description(self, event_id: str, source: str) -> None: try: event: Event = Event.get(Event.id == event_id) except DoesNotExist: @@ -256,4 +258,23 @@ class EmbeddingMaintainer(threading.Thread): metadata = get_metadata(event) thumbnail = base64.b64decode(event.thumbnail) - self._embed_description(event, [thumbnail], metadata) + logger.debug(f"Using ${source} regeneration for ${event}") + + if event.has_snapshot and source == "snapshot": + with open( + os.path.join(CLIPS_DIR, f"{event.camera}-{event.id}.jpg"), + "rb", + ) as image_file: + snapshot_image = image_file.read() + + embed_image = ( + [snapshot_image] + if event.has_snapshot and source == "snapshot" + else ( + [thumbnail for data in self.tracked_events[event_id]] + if len(self.tracked_events.get(event_id, [])) > 0 + else [thumbnail] + ) + ) + + self._embed_description(event, embed_image, metadata) diff --git a/frigate/events/types.py b/frigate/events/types.py index 1750b3e7b..1461c1f28 100644 --- a/frigate/events/types.py +++ b/frigate/events/types.py @@ -12,3 +12,8 @@ class EventStateEnum(str, Enum): start = "start" update = "update" end = "end" + + +class RegenerateDescriptionEnum(str, Enum): + thumbnails = "thumbnails" + snapshot = "snapshot"