From a10a49a85c82c988a3be7196be6bcce6d09855e4 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Mon, 21 Oct 2024 11:04:57 -0600
Subject: [PATCH] Implement face embedding calculation

---
 frigate/db/sqlitevecq.py         |  3 +--
 frigate/embeddings/embeddings.py | 22 ++++++++++++++++++++++
 frigate/embeddings/maintainer.py |  7 ++++++-
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/frigate/db/sqlitevecq.py b/frigate/db/sqlitevecq.py
index cd63c1b7a..b852e06e5 100644
--- a/frigate/db/sqlitevecq.py
+++ b/frigate/db/sqlitevecq.py
@@ -59,7 +59,6 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase):
             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 5586b4659..c7b4c3d08 100644
--- a/frigate/embeddings/embeddings.py
+++ b/frigate/embeddings/embeddings.py
@@ -3,6 +3,8 @@
 import base64
 import logging
 import os
+import random
+import string
 import time
 
 from numpy import ndarray
@@ -215,6 +217,26 @@ class Embeddings:
 
         return embeddings
 
+    def embed_face(self, label: str, thumbnail: bytes, upsert: bool = False) -> ndarray:
+        # Convert thumbnail bytes to PIL Image
+        image = Image.open(io.BytesIO(thumbnail)).convert("RGB")
+        embedding = self.vision_embedding([image])[0]
+
+        if upsert:
+            rand_id = "".join(
+                random.choices(string.ascii_lowercase + string.digits, k=6)
+            )
+            id = f"{label}-{rand_id}"
+            self.db.execute_sql(
+                """
+                INSERT OR REPLACE INTO vec_faces(id, face_embedding)
+                VALUES(?, ?)
+                """,
+                (id, serialize(embedding)),
+            )
+
+        return embedding
+
     def reindex(self) -> None:
         logger.info("Indexing tracked object embeddings...")
 
diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py
index dee78fd0a..68a0dc763 100644
--- a/frigate/embeddings/maintainer.py
+++ b/frigate/embeddings/maintainer.py
@@ -295,6 +295,10 @@ class EmbeddingMaintainer(threading.Thread):
 
         face_frame = cv2.cvtColor(frame, cv2.COLOR_YUV2BGR_I420)
         face_frame = face_frame[face_box[1] : face_box[3], face_box[0] : face_box[2]]
+        ret, jpg = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
+
+        if not ret:
+            return
 
         # if face_frame is not None:
         #    cv2.imwrite(
@@ -303,7 +307,8 @@ class EmbeddingMaintainer(threading.Thread):
         #        [int(cv2.IMWRITE_WEBP_QUALITY), 60],
         #    )
 
-        # TODO run embedding on face box
+        embedding = self.embeddings.embed_face("nick", jpg.tobytes(), upsert=False)
+
         # TODO compare embedding to faces in embeddings DB to fine cosine similarity
         # TODO check against threshold and min score to see if best face qualifies
         # TODO update tracked object