administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-05-05 13:07:44 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Rewrite build model task to be asynchronous so it doesn't block the pipeline

Browse Source
This commit is contained in:
Nicolas Mowen 2025-03-26 09:13:04 -06:00
parent 2b1fc069eb
commit 900f148aa0
1 changed files with 51 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
frigate/data_processing/common/face/model.py
View File

@ -1,5 +1,7 @@
import logging
import os
import queue
import threading
from abc import ABC, abstractmethod
import cv2
@ -206,46 +208,69 @@ class ArcFaceRecognizer(FaceRecognizer):
super().__init__(config)
self.mean_embs: dict[int, np.ndarray] = {}
self.face_embedder: ArcfaceEmbedding = ArcfaceEmbedding()
self.model_builder_queue: queue.Queue | None = None
def clear(self) -> None:
self.mean_embs = {}
def run_build_task(self) -> None:
self.model_builder_queue = queue.Queue()
def build_model():
face_embeddings_map: dict[str, list[np.ndarray]] = {}
idx = 0
dir = "/media/frigate/clips/faces"
for name in os.listdir(dir):
if name == "train":
continue
face_folder = os.path.join(dir, name)
if not os.path.isdir(face_folder):
continue
face_embeddings_map[name] = []
for image in os.listdir(face_folder):
img = cv2.imread(os.path.join(face_folder, image))
if img is None:
continue
img = self.align_face(img, img.shape[1], img.shape[0])
emb = self.face_embedder([img])[0].squeeze()
face_embeddings_map[name].append(emb)
idx += 1
self.model_builder_queue.put(face_embeddings_map)
thread = threading.Thread(target=build_model)
thread.start()
def build(self):
if not self.landmark_detector:
self.init_landmark_detector()
return None
face_embeddings_map: dict[str, list[np.ndarray]] = {}
idx = 0
dir = "/media/frigate/clips/faces"
for name in os.listdir(dir):
if name == "train":
continue
face_folder = os.path.join(dir, name)
if not os.path.isdir(face_folder):
continue
face_embeddings_map[name] = []
for image in os.listdir(face_folder):
img = cv2.imread(os.path.join(face_folder, image))
if img is None:
continue
img = self.align_face(img, img.shape[1], img.shape[0])
emb = self.face_embedder([img])[0].squeeze()
face_embeddings_map[name].append(emb)
idx += 1
if self.model_builder_queue is not None:
try:
face_embeddings_map: dict[str, list[np.ndarray]] = (
self.model_builder_queue.get(timeout=0.1)
)
self.model_builder_queue = None
except queue.Empty:
return
else:
self.run_build_task()
return
if not face_embeddings_map:
return
for name, embs in face_embeddings_map.items():
self.mean_embs[name] = stats.trim_mean(embs, 0.15)
if embs:
self.mean_embs[name] = stats.trim_mean(embs, 0.15)
logger.debug("Finished building ArcFace model")