mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-05-05 13:07:44 +03:00
Rewrite build model task to be asynchronous so it doesn't block the pipeline
This commit is contained in:
parent
2b1fc069eb
commit
900f148aa0
@ -1,5 +1,7 @@
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
|
import queue
|
||||||
|
import threading
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
@ -206,15 +208,15 @@ class ArcFaceRecognizer(FaceRecognizer):
|
|||||||
super().__init__(config)
|
super().__init__(config)
|
||||||
self.mean_embs: dict[int, np.ndarray] = {}
|
self.mean_embs: dict[int, np.ndarray] = {}
|
||||||
self.face_embedder: ArcfaceEmbedding = ArcfaceEmbedding()
|
self.face_embedder: ArcfaceEmbedding = ArcfaceEmbedding()
|
||||||
|
self.model_builder_queue: queue.Queue | None = None
|
||||||
|
|
||||||
def clear(self) -> None:
|
def clear(self) -> None:
|
||||||
self.mean_embs = {}
|
self.mean_embs = {}
|
||||||
|
|
||||||
def build(self):
|
def run_build_task(self) -> None:
|
||||||
if not self.landmark_detector:
|
self.model_builder_queue = queue.Queue()
|
||||||
self.init_landmark_detector()
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
def build_model():
|
||||||
face_embeddings_map: dict[str, list[np.ndarray]] = {}
|
face_embeddings_map: dict[str, list[np.ndarray]] = {}
|
||||||
idx = 0
|
idx = 0
|
||||||
|
|
||||||
@ -241,10 +243,33 @@ class ArcFaceRecognizer(FaceRecognizer):
|
|||||||
|
|
||||||
idx += 1
|
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
|
||||||
|
|
||||||
|
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:
|
if not face_embeddings_map:
|
||||||
return
|
return
|
||||||
|
|
||||||
for name, embs in face_embeddings_map.items():
|
for name, embs in face_embeddings_map.items():
|
||||||
|
if embs:
|
||||||
self.mean_embs[name] = stats.trim_mean(embs, 0.15)
|
self.mean_embs[name] = stats.trim_mean(embs, 0.15)
|
||||||
|
|
||||||
logger.debug("Finished building ArcFace model")
|
logger.debug("Finished building ArcFace model")
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user