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

This commit is contained in:
Nicolas Mowen 2025-03-26 09:13:04 -06:00
parent 2b1fc069eb
commit 900f148aa0

View File

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