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
__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."
)

View File

@ -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
);
""")

View File

@ -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