diff --git a/frigate/config/semantic_search.py b/frigate/config/semantic_search.py index 2891050a1..9ed06292a 100644 --- a/frigate/config/semantic_search.py +++ b/frigate/config/semantic_search.py @@ -4,7 +4,17 @@ from pydantic import Field from .base import FrigateBaseModel -__all__ = ["SemanticSearchConfig"] +__all__ = ["FaceRecognitionConfig", "SemanticSearchConfig"] + + +class FaceRecognitionConfig(FrigateBaseModel): + enabled: bool = Field(default=False, title="Enable face recognition.") + reindex: Optional[bool] = Field( + default=False, title="Reindex all detections on startup." + ) + model_size: str = Field( + default="small", title="The size of the embeddings model used." + ) class SemanticSearchConfig(FrigateBaseModel): @@ -12,6 +22,9 @@ class SemanticSearchConfig(FrigateBaseModel): reindex: Optional[bool] = Field( default=False, title="Reindex all detections on startup." ) + face_recognition: FaceRecognitionConfig = Field( + default_factory=FaceRecognitionConfig, title="Face recognition config." + ) model_size: str = Field( default="small", title="The size of the embeddings model used." ) diff --git a/frigate/db/sqlitevecq.py b/frigate/db/sqlitevecq.py index ccb75ae54..cd63c1b7a 100644 --- a/frigate/db/sqlitevecq.py +++ b/frigate/db/sqlitevecq.py @@ -36,8 +36,11 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase): self.execute_sql(""" DROP TABLE vec_thumbnails; """) + self.execute_sql(""" + DROP TABLE vec_faces; + """) - def create_embeddings_tables(self) -> None: + def create_embeddings_tables(self, face_recognition: bool) -> None: """Create vec0 virtual table for embeddings""" self.execute_sql(""" CREATE VIRTUAL TABLE IF NOT EXISTS vec_thumbnails USING vec0( @@ -51,3 +54,12 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase): description_embedding FLOAT[768] distance_metric=cosine ); """) + + if face_recognition: + self.execute_sql(""" + CREATE VIRTUAL TABLE IF NOT EXISTS vec_faces USING vec0( + id TEXT PRIMARY KEY, + faceName TEXT, + face_embedding FLOAT[768] distance_metric=cosine + ); + """) \ No newline at end of file diff --git a/frigate/embeddings/embeddings.py b/frigate/embeddings/embeddings.py index d77a9eecf..5586b4659 100644 --- a/frigate/embeddings/embeddings.py +++ b/frigate/embeddings/embeddings.py @@ -67,7 +67,7 @@ class Embeddings: self.requestor = InterProcessRequestor() # Create tables if they don't exist - self.db.create_embeddings_tables() + self.db.create_embeddings_tables(self.config.face_recognition.enabled) models = [ "jinaai/jina-clip-v1-text_model_fp16.onnx", @@ -220,7 +220,7 @@ class Embeddings: self.db.drop_embeddings_tables() logger.debug("Dropped embeddings tables.") - self.db.create_embeddings_tables() + self.db.create_embeddings_tables(self.config.face_recognition.enabled) logger.debug("Created embeddings tables.") # Delete the saved stats file