From f447087d4af4c90b27c27751a7466577d9773df6 Mon Sep 17 00:00:00 2001
From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com>
Date: Mon, 30 Sep 2024 12:29:41 -0500
Subject: [PATCH] Allow embedding of snapshot for description via config option

---
 frigate/config/camera/genai.py   |  3 +++
 frigate/embeddings/maintainer.py | 29 +++++++++++++++++++++--------
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/frigate/config/camera/genai.py b/frigate/config/camera/genai.py
index 736974535..12ff37adb 100644
--- a/frigate/config/camera/genai.py
+++ b/frigate/config/camera/genai.py
@@ -18,6 +18,9 @@ class GenAIProviderEnum(str, Enum):
 # uses BaseModel because some global attributes are not available at the camera level
 class GenAICameraConfig(BaseModel):
     enabled: bool = Field(default=False, title="Enable GenAI for camera.")
+    use_snapshot: bool = Field(
+        default=False, title="Use snapshots for generating descriptions."
+    )
     prompt: str = Field(
         default="Describe the {label} in the sequence of images with as much detail as possible. Do not describe the background.",
         title="Default caption prompt.",
diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py
index cbe4554ce..31f5bbd27 100644
--- a/frigate/embeddings/maintainer.py
+++ b/frigate/embeddings/maintainer.py
@@ -3,6 +3,7 @@
 import base64
 import io
 import logging
+import os
 import threading
 from multiprocessing.synchronize import Event as MpEvent
 from typing import Optional
@@ -19,7 +20,7 @@ from frigate.comms.event_metadata_updater import (
 from frigate.comms.events_updater import EventEndSubscriber, EventUpdateSubscriber
 from frigate.comms.inter_process import InterProcessRequestor
 from frigate.config import FrigateConfig
-from frigate.const import UPDATE_EVENT_DESCRIPTION
+from frigate.const import CLIPS_DIR, UPDATE_EVENT_DESCRIPTION
 from frigate.events.types import EventTypeEnum
 from frigate.genai import get_genai_client
 from frigate.models import Event
@@ -136,19 +137,31 @@ class EmbeddingMaintainer(threading.Thread):
                         or set(event.zones) & set(camera_config.genai.required_zones)
                     )
                 ):
-                    # Generate the description. Call happens in a thread since it is network bound.
+                    if event.has_snapshot and camera_config.genai.use_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 camera_config.genai.use_snapshot
+                        else (
+                            [thumbnail for data in self.tracked_events[event_id]]
+                            if len(self.tracked_events.get(event_id, [])) > 0
+                            else [thumbnail]
+                        )
+                    )
+
+                    # Generate the description, passing the correct image argument
                     threading.Thread(
                         target=self._embed_description,
                         name=f"_embed_description_{event.id}",
                         daemon=True,
                         args=(
                             event,
-                            [
-                                data["thumbnail"]
-                                for data in self.tracked_events[event_id]
-                            ]
-                            if len(self.tracked_events.get(event_id, [])) > 0
-                            else [thumbnail],
+                            embed_image,
                             metadata,
                         ),
                     ).start()