Face model loading improvements (#17390)

* Don't assume landmark file is downloaded

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

* Handle case where face recognition does not respond

* Cleanup

* Make daemon thread
This commit is contained in:
Nicolas Mowen 2025-03-26 10:00:23 -06:00 committed by GitHub
parent e3d4b84803
commit e6936c177b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 80 additions and 29 deletions

View File

@ -198,6 +198,16 @@ async def register_face(request: Request, name: str, file: UploadFile):
context: EmbeddingsContext = request.app.embeddings
result = context.register_face(name, await file.read())
if not isinstance(result, dict):
return JSONResponse(
status_code=500,
content={
"success": False,
"message": "Could not process request. Try restarting Frigate.",
},
)
return JSONResponse(
status_code=200 if result.get("success", True) else 400,
content=result,
@ -214,6 +224,16 @@ async def recognize_face(request: Request, file: UploadFile):
context: EmbeddingsContext = request.app.embeddings
result = context.recognize_face(await file.read())
if not isinstance(result, dict):
return JSONResponse(
status_code=500,
content={
"success": False,
"message": "Could not process request. Try restarting Frigate.",
},
)
return JSONResponse(
status_code=200 if result.get("success", True) else 400,
content=result,

View File

@ -1,5 +1,7 @@
import logging
import os
import queue
import threading
from abc import ABC, abstractmethod
import cv2
@ -18,10 +20,7 @@ class FaceRecognizer(ABC):
def __init__(self, config: FrigateConfig) -> None:
self.config = config
self.landmark_detector = cv2.face.createFacemarkLBF()
self.landmark_detector.loadModel(
os.path.join(MODEL_CACHE_DIR, "facedet/landmarkdet.yaml")
)
self.init_landmark_detector()
@abstractmethod
def build(self) -> None:
@ -37,6 +36,13 @@ class FaceRecognizer(ABC):
def classify(self, face_image: np.ndarray) -> tuple[str, float] | None:
pass
def init_landmark_detector(self) -> None:
landmark_model = os.path.join(MODEL_CACHE_DIR, "facedet/landmarkdet.yaml")
if os.path.exists(landmark_model):
self.landmark_detector = cv2.face.createFacemarkLBF()
self.landmark_detector.loadModel(landmark_model)
def align_face(
self,
image: np.ndarray,
@ -130,6 +136,7 @@ class LBPHRecognizer(FaceRecognizer):
def build(self):
if not self.landmark_detector:
self.init_landmark_detector()
return None
labels = []
@ -201,45 +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 build(self):
if not self.landmark_detector:
return None
def run_build_task(self) -> None:
self.model_builder_queue = queue.Queue()
face_embeddings_map: dict[str, list[np.ndarray]] = {}
idx = 0
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:
dir = "/media/frigate/clips/faces"
for name in os.listdir(dir):
if name == "train":
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)
face_folder = os.path.join(dir, name)
idx += 1
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, daemon=True)
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:
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")