Add basic config and face recognition table

This commit is contained in:
Nicolas Mowen 2024-10-21 10:21:17 -06:00
parent 40c6fda19d
commit 819e6e4d40
3 changed files with 29 additions and 4 deletions

View File

@ -4,7 +4,17 @@ from pydantic import Field
from .base import FrigateBaseModel 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): class SemanticSearchConfig(FrigateBaseModel):
@ -12,6 +22,9 @@ class SemanticSearchConfig(FrigateBaseModel):
reindex: Optional[bool] = Field( reindex: Optional[bool] = Field(
default=False, title="Reindex all detections on startup." default=False, title="Reindex all detections on startup."
) )
face_recognition: FaceRecognitionConfig = Field(
default_factory=FaceRecognitionConfig, title="Face recognition config."
)
model_size: str = Field( model_size: str = Field(
default="small", title="The size of the embeddings model used." default="small", title="The size of the embeddings model used."
) )

View File

@ -36,8 +36,11 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase):
self.execute_sql(""" self.execute_sql("""
DROP TABLE vec_thumbnails; 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""" """Create vec0 virtual table for embeddings"""
self.execute_sql(""" self.execute_sql("""
CREATE VIRTUAL TABLE IF NOT EXISTS vec_thumbnails USING vec0( CREATE VIRTUAL TABLE IF NOT EXISTS vec_thumbnails USING vec0(
@ -51,3 +54,12 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase):
description_embedding FLOAT[768] distance_metric=cosine 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
);
""")

View File

@ -67,7 +67,7 @@ class Embeddings:
self.requestor = InterProcessRequestor() self.requestor = InterProcessRequestor()
# Create tables if they don't exist # Create tables if they don't exist
self.db.create_embeddings_tables() self.db.create_embeddings_tables(self.config.face_recognition.enabled)
models = [ models = [
"jinaai/jina-clip-v1-text_model_fp16.onnx", "jinaai/jina-clip-v1-text_model_fp16.onnx",
@ -220,7 +220,7 @@ class Embeddings:
self.db.drop_embeddings_tables() self.db.drop_embeddings_tables()
logger.debug("Dropped 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.") logger.debug("Created embeddings tables.")
# Delete the saved stats file # Delete the saved stats file