blakeblackshear.frigate/frigate/embeddings/__init__.py

"""SQLite-vec embeddings database."""

import base64
import json
import logging
import os
import threading
from multiprocessing.synchronize import Event as MpEvent
from typing import Any, Union

import regex
from pathvalidate import ValidationError, sanitize_filename

from frigate.comms.embeddings_updater import EmbeddingsRequestEnum, EmbeddingsRequestor
from frigate.config import FrigateConfig
from frigate.const import CONFIG_DIR, FACE_DIR, PROCESS_PRIORITY_HIGH
from frigate.data_processing.types import DataProcessorMetrics
from frigate.db.sqlitevecq import SqliteVecQueueDatabase
from frigate.models import Event
from frigate.util.builtin import serialize
from frigate.util.classification import kickoff_model_training
from frigate.util.process import FrigateProcess

from .maintainer import EmbeddingMaintainer
from .util import ZScoreNormalization

logger = logging.getLogger(__name__)


class EmbeddingProcess(FrigateProcess):
    def __init__(
        self,
        config: FrigateConfig,
        metrics: DataProcessorMetrics | None,
        stop_event: MpEvent,
    ) -> None:
        super().__init__(
            stop_event,
            PROCESS_PRIORITY_HIGH,
            name="frigate.embeddings_manager",
            daemon=True,
        )
        self.config = config
        self.metrics = metrics

    def run(self) -> None:
        self.pre_run_setup(self.config.logger)
        maintainer = EmbeddingMaintainer(
            self.config,
            self.metrics,
            self.stop_event,
        )
        maintainer.start()


class EmbeddingsContext:
    def __init__(self, db: SqliteVecQueueDatabase):
        self.db = db
        self.thumb_stats = ZScoreNormalization()
        self.desc_stats = ZScoreNormalization()
        self.requestor = EmbeddingsRequestor()

        # load stats from disk
        try:
            with open(os.path.join(CONFIG_DIR, ".search_stats.json"), "r") as f:
                data = json.loads(f.read())
                self.thumb_stats.from_dict(data["thumb_stats"])
                self.desc_stats.from_dict(data["desc_stats"])
        except FileNotFoundError:
            pass

    def stop(self):
        """Write the stats to disk as JSON on exit."""
        contents = {
            "thumb_stats": self.thumb_stats.to_dict(),
            "desc_stats": self.desc_stats.to_dict(),
        }
        with open(os.path.join(CONFIG_DIR, ".search_stats.json"), "w") as f:
            json.dump(contents, f)
        self.requestor.stop()

    def search_thumbnail(
        self, query: Union[Event, str], event_ids: list[str] = None
    ) -> list[tuple[str, float]]:
        if query.__class__ == Event:
            cursor = self.db.execute_sql(
                """
                SELECT thumbnail_embedding FROM vec_thumbnails WHERE id = ?
                """,
                [query.id],
            )

            row = cursor.fetchone() if cursor else None

            if row:
                query_embedding = row[0]
            else:
                # If no embedding found, generate it and return it
                data = self.requestor.send_data(
                    EmbeddingsRequestEnum.embed_thumbnail.value,
                    {"id": str(query.id), "thumbnail": str(query.thumbnail)},
                )

                if not data:
                    return []

                query_embedding = serialize(data)
        else:
            data = self.requestor.send_data(
                EmbeddingsRequestEnum.generate_search.value, query
            )

            if not data:
                return []

            query_embedding = serialize(data)

        sql_query = """
            SELECT
                id,
                distance
            FROM vec_thumbnails
            WHERE thumbnail_embedding MATCH ?
                AND k = 100
        """

        # Add the IN clause if event_ids is provided and not empty
        # this is the only filter supported by sqlite-vec as of 0.1.3
        # but it seems to be broken in this version
        if event_ids:
            sql_query += " AND id IN ({})".format(",".join("?" * len(event_ids)))

        # order by distance DESC is not implemented in this version of sqlite-vec
        # when it's implemented, we can use cosine similarity
        sql_query += " ORDER BY distance"

        parameters = [query_embedding] + event_ids if event_ids else [query_embedding]

        results = self.db.execute_sql(sql_query, parameters).fetchall()

        return results

    def search_description(
        self, query_text: str, event_ids: list[str] = None
    ) -> list[tuple[str, float]]:
        data = self.requestor.send_data(
            EmbeddingsRequestEnum.generate_search.value, query_text
        )

        if not data:
            return []

        query_embedding = serialize(data)

        # Prepare the base SQL query
        sql_query = """
            SELECT
                id,
                distance
            FROM vec_descriptions
            WHERE description_embedding MATCH ?
                AND k = 100
        """

        # Add the IN clause if event_ids is provided and not empty
        # this is the only filter supported by sqlite-vec as of 0.1.3
        # but it seems to be broken in this version
        if event_ids:
            sql_query += " AND id IN ({})".format(",".join("?" * len(event_ids)))

        # order by distance DESC is not implemented in this version of sqlite-vec
        # when it's implemented, we can use cosine similarity
        sql_query += " ORDER BY distance"

        parameters = [query_embedding] + event_ids if event_ids else [query_embedding]

        results = self.db.execute_sql(sql_query, parameters).fetchall()

        return results

    def register_face(self, face_name: str, image_data: bytes) -> dict[str, Any]:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.register_face.value,
            {
                "face_name": face_name,
                "image": base64.b64encode(image_data).decode("ASCII"),
            },
        )

    def recognize_face(self, image_data: bytes) -> dict[str, Any]:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.recognize_face.value,
            {
                "image": base64.b64encode(image_data).decode("ASCII"),
            },
        )

    def get_face_ids(self, name: str) -> list[str]:
        sql_query = f"""
            SELECT
                id
            FROM vec_descriptions
            WHERE id LIKE '%{name}%'
        """

        return self.db.execute_sql(sql_query).fetchall()

    def reprocess_face(self, face_file: str) -> dict[str, Any]:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.reprocess_face.value, {"image_file": face_file}
        )

    def clear_face_classifier(self) -> None:
        self.requestor.send_data(
            EmbeddingsRequestEnum.clear_face_classifier.value, None
        )

    def delete_face_ids(self, face: str, ids: list[str]) -> None:
        folder = os.path.join(FACE_DIR, face)
        for id in ids:
            file_path = os.path.join(folder, id)

            if os.path.isfile(file_path):
                os.unlink(file_path)

        if face != "train" and len(os.listdir(folder)) == 0:
            os.rmdir(folder)

        self.requestor.send_data(
            EmbeddingsRequestEnum.clear_face_classifier.value, None
        )

    def rename_face(self, old_name: str, new_name: str) -> None:
        valid_name_pattern = r"^[\p{L}\p{N}\s'_-]{1,50}$"

        try:
            sanitized_old_name = sanitize_filename(old_name, replacement_text="_")
            sanitized_new_name = sanitize_filename(new_name, replacement_text="_")
        except ValidationError as e:
            raise ValueError(f"Invalid face name: {str(e)}")

        if not regex.match(valid_name_pattern, old_name):
            raise ValueError(f"Invalid old face name: {old_name}")
        if not regex.match(valid_name_pattern, new_name):
            raise ValueError(f"Invalid new face name: {new_name}")
        if sanitized_old_name != old_name:
            raise ValueError(f"Old face name contains invalid characters: {old_name}")
        if sanitized_new_name != new_name:
            raise ValueError(f"New face name contains invalid characters: {new_name}")

        old_path = os.path.normpath(os.path.join(FACE_DIR, old_name))
        new_path = os.path.normpath(os.path.join(FACE_DIR, new_name))

        # Prevent path traversal
        if not old_path.startswith(
            os.path.normpath(FACE_DIR)
        ) or not new_path.startswith(os.path.normpath(FACE_DIR)):
            raise ValueError("Invalid path detected")

        if not os.path.exists(old_path):
            raise ValueError(f"Face {old_name} not found.")

        os.rename(old_path, new_path)

        self.requestor.send_data(
            EmbeddingsRequestEnum.clear_face_classifier.value, None
        )

    def update_description(self, event_id: str, description: str) -> None:
        self.requestor.send_data(
            EmbeddingsRequestEnum.embed_description.value,
            {"id": event_id, "description": description},
        )

    def reprocess_plate(self, event: dict[str, Any]) -> dict[str, Any]:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.reprocess_plate.value, {"event": event}
        )

    def reindex_embeddings(self) -> dict[str, Any]:
        return self.requestor.send_data(EmbeddingsRequestEnum.reindex.value, {})

    def start_classification_training(self, model_name: str) -> dict[str, Any]:
        threading.Thread(
            target=kickoff_model_training,
            args=(self.requestor, model_name),
            daemon=True,
        ).start()
        return {"success": True, "message": f"Began training {model_name} model."}

    def transcribe_audio(self, event: dict[str, any]) -> dict[str, any]:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.transcribe_audio.value, {"event": event}
        )

    def generate_description_embedding(self, text: str) -> None:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.embed_description.value,
            {"id": None, "description": text, "upsert": False},
        )

    def generate_image_embedding(self, event_id: str, thumbnail: bytes) -> None:
        return self.requestor.send_data(
            EmbeddingsRequestEnum.embed_thumbnail.value,
            {"id": str(event_id), "thumbnail": str(thumbnail), "upsert": False},
        )