mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-02-15 15:45:27 +03:00
Implement face embedding calculation
This commit is contained in:
parent
33b0958531
commit
a10a49a85c
@ -59,7 +59,6 @@ class SqliteVecQueueDatabase(SqliteQueueDatabase):
|
|||||||
self.execute_sql("""
|
self.execute_sql("""
|
||||||
CREATE VIRTUAL TABLE IF NOT EXISTS vec_faces USING vec0(
|
CREATE VIRTUAL TABLE IF NOT EXISTS vec_faces USING vec0(
|
||||||
id TEXT PRIMARY KEY,
|
id TEXT PRIMARY KEY,
|
||||||
faceName TEXT,
|
|
||||||
face_embedding FLOAT[768] distance_metric=cosine
|
face_embedding FLOAT[768] distance_metric=cosine
|
||||||
);
|
);
|
||||||
""")
|
""")
|
||||||
|
|||||||
@ -3,6 +3,8 @@
|
|||||||
import base64
|
import base64
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
|
import random
|
||||||
|
import string
|
||||||
import time
|
import time
|
||||||
|
|
||||||
from numpy import ndarray
|
from numpy import ndarray
|
||||||
@ -215,6 +217,26 @@ class Embeddings:
|
|||||||
|
|
||||||
return 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:
|
def reindex(self) -> None:
|
||||||
logger.info("Indexing tracked object embeddings...")
|
logger.info("Indexing tracked object embeddings...")
|
||||||
|
|
||||||
|
|||||||
@ -295,6 +295,10 @@ class EmbeddingMaintainer(threading.Thread):
|
|||||||
|
|
||||||
face_frame = cv2.cvtColor(frame, cv2.COLOR_YUV2BGR_I420)
|
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]]
|
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:
|
# if face_frame is not None:
|
||||||
# cv2.imwrite(
|
# cv2.imwrite(
|
||||||
@ -303,7 +307,8 @@ class EmbeddingMaintainer(threading.Thread):
|
|||||||
# [int(cv2.IMWRITE_WEBP_QUALITY), 60],
|
# [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 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 check against threshold and min score to see if best face qualifies
|
||||||
# TODO update tracked object
|
# TODO update tracked object
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user