From 2c182d59f02ed8af872586c2615c908f039d4414 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Fri, 3 Jan 2025 09:23:30 -0700 Subject: [PATCH] Support downloading face models --- frigate/embeddings/embeddings.py | 13 ------ frigate/embeddings/maintainer.py | 18 +-------- frigate/util/model.py | 68 ++++++++++++++++++++++++++------ 3 files changed, 59 insertions(+), 40 deletions(-) diff --git a/frigate/embeddings/embeddings.py b/frigate/embeddings/embeddings.py index f35158aef..63597e49e 100644 --- a/frigate/embeddings/embeddings.py +++ b/frigate/embeddings/embeddings.py @@ -123,19 +123,6 @@ class Embeddings: device="GPU" if config.semantic_search.model_size == "large" else "CPU", ) - if self.config.face_recognition.enabled: - self.face_embedding = GenericONNXEmbedding( - model_name="facedet", - model_file="facedet.onnx", - download_urls={ - "facedet.onnx": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/facedet.onnx", - "landmarkdet.yaml": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/landmarkdet.yaml", - }, - model_size="small", - model_type=ModelTypeEnum.face, - requestor=self.requestor, - ) - self.lpr_detection_model = None self.lpr_classification_model = None self.lpr_recognition_model = None diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py index b5c050325..175b8d4e9 100644 --- a/frigate/embeddings/maintainer.py +++ b/frigate/embeddings/maintainer.py @@ -100,19 +100,6 @@ class EmbeddingMaintainer(threading.Thread): self.lpr_config, self.requestor, self.embeddings ) - @property - def face_detector(self) -> cv2.FaceDetectorYN: - # Lazily create the classifier. - if "face_detector" not in self.__dict__: - self.__dict__["face_detector"] = cv2.FaceDetectorYN.create( - "/config/model_cache/facedet/facedet.onnx", - config="", - input_size=(320, 320), - score_threshold=0.8, - nms_threshold=0.3, - ) - return self.__dict__["face_detector"] - def run(self) -> None: """Maintain a SQLite-vec database for semantic search.""" while not self.stop_event.is_set(): @@ -395,10 +382,9 @@ class EmbeddingMaintainer(threading.Thread): def _detect_face(self, input: np.ndarray) -> tuple[int, int, int, int]: """Detect faces in input image.""" - self.face_detector.setInputSize((input.shape[1], input.shape[0])) - faces = self.face_detector.detect(input) + faces = self.face_classifier.detect_faces(input) - if faces[1] is None: + if faces is None or faces[1] is None: return None face = None diff --git a/frigate/util/model.py b/frigate/util/model.py index 19c4b4e76..106815b96 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -2,7 +2,7 @@ import logging import os -from typing import Any, Optional +from typing import Any import cv2 import numpy as np @@ -10,6 +10,8 @@ import onnxruntime as ort from playhouse.sqliteq import SqliteQueueDatabase from frigate.config.semantic_search import FaceRecognitionConfig +from frigate.const import MODEL_CACHE_DIR +from frigate.util.downloader import ModelDownloader try: import openvino as ov @@ -162,21 +164,50 @@ class FaceClassificationModel: def __init__(self, config: FaceRecognitionConfig, db: SqliteQueueDatabase): self.config = config self.db = db - self.landmark_detector = cv2.face.createFacemarkLBF() + self.face_detector: cv2.FaceDetectorYN = None + self.landmark_detector: cv2.face.FacemarkLBF = None + self.face_recognizer: cv2.face.LBPHFaceRecognizer = None - if os.path.isfile("/config/model_cache/facedet/landmarkdet.yaml"): - self.landmark_detector.loadModel( - "/config/model_cache/facedet/landmarkdet.yaml" - ) + download_path = os.path.join(MODEL_CACHE_DIR, "facedet") + model_files = ["facedet.onnx", "landmarkdet.yaml"] + self.model_urls = [ + "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/facedet.onnx", + "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/landmarkdet.yaml", + ] - self.recognizer: cv2.face.LBPHFaceRecognizer = ( - cv2.face.LBPHFaceRecognizer_create( - radius=2, threshold=(1 - config.min_score) * 1000 + if not all(os.path.exists(os.path.join(download_path, n)) for n in model_files): + logger.debug(f"starting model download for {self.model_name}") + self.downloader = ModelDownloader( + model_name="facedet.onnx", + download_path=download_path, + file_names=model_files, + download_func=self.__download_models, ) - ) + self.downloader.ensure_model_files() + else: + self.__build_detector() + self.label_map: dict[int, str] = {} self.__build_classifier() + def __download_models(self, path: str) -> None: + try: + file_name = os.path.basename(path) + ModelDownloader.download_from_url(self.model_urls[file_name], path) + except Exception: + pass + + def __build_detector(self) -> None: + self.face_detector = cv2.FaceDetectorYN.create( + "/config/model_cache/facedet/facedet.onnx", + config="", + input_size=(320, 320), + score_threshold=0.8, + nms_threshold=0.3, + ) + self.landmark_detector = cv2.face.createFacemarkLBF() + self.landmark_detector.loadModel("/config/model_cache/facedet/landmarkdet.yaml") + def __build_classifier(self) -> None: labels = [] faces = [] @@ -203,6 +234,11 @@ class FaceClassificationModel: faces.append(img) labels.append(idx) + self.recognizer: cv2.face.LBPHFaceRecognizer = ( + cv2.face.LBPHFaceRecognizer_create( + radius=2, threshold=(1 - self.config.min_score) * 1000 + ) + ) self.recognizer.train(faces, np.array(labels)) def __align_face( @@ -267,7 +303,17 @@ class FaceClassificationModel: self.labeler = None self.label_map = {} - def classify_face(self, face_image: np.ndarray) -> Optional[tuple[str, float]]: + def detect_faces(self, input: np.ndarray) -> tuple[int, cv2.typing.MatLike] | None: + if not self.face_detector: + return None + + self.face_detector.setInputSize((input.shape[1], input.shape[0])) + return self.face_detector.detect(input) + + def classify_face(self, face_image: np.ndarray) -> tuple[str, float] | None: + if not self.landmark_detector: + return None + if not self.label_map: self.__build_classifier()