Allow embedding of snapshot for description via config option

This commit is contained in:
Josh Hawkins 2024-09-30 12:29:41 -05:00
parent 3688a3bc67
commit f447087d4a
2 changed files with 24 additions and 8 deletions

View File

@ -18,6 +18,9 @@ class GenAIProviderEnum(str, Enum):
# uses BaseModel because some global attributes are not available at the camera level # uses BaseModel because some global attributes are not available at the camera level
class GenAICameraConfig(BaseModel): class GenAICameraConfig(BaseModel):
enabled: bool = Field(default=False, title="Enable GenAI for camera.") 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( prompt: str = Field(
default="Describe the {label} in the sequence of images with as much detail as possible. Do not describe the background.", default="Describe the {label} in the sequence of images with as much detail as possible. Do not describe the background.",
title="Default caption prompt.", title="Default caption prompt.",

View File

@ -3,6 +3,7 @@
import base64 import base64
import io import io
import logging import logging
import os
import threading import threading
from multiprocessing.synchronize import Event as MpEvent from multiprocessing.synchronize import Event as MpEvent
from typing import Optional 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.events_updater import EventEndSubscriber, EventUpdateSubscriber
from frigate.comms.inter_process import InterProcessRequestor from frigate.comms.inter_process import InterProcessRequestor
from frigate.config import FrigateConfig 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.events.types import EventTypeEnum
from frigate.genai import get_genai_client from frigate.genai import get_genai_client
from frigate.models import Event from frigate.models import Event
@ -136,19 +137,31 @@ class EmbeddingMaintainer(threading.Thread):
or set(event.zones) & set(camera_config.genai.required_zones) 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( threading.Thread(
target=self._embed_description, target=self._embed_description,
name=f"_embed_description_{event.id}", name=f"_embed_description_{event.id}",
daemon=True, daemon=True,
args=( args=(
event, event,
[ embed_image,
data["thumbnail"]
for data in self.tracked_events[event_id]
]
if len(self.tracked_events.get(event_id, [])) > 0
else [thumbnail],
metadata, metadata,
), ),
).start() ).start()