From cc03075cbecab66ce8f539f30c0067708dc3fcf9 Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Tue, 25 Mar 2025 14:04:41 -0600 Subject: [PATCH] Upgrade to new arcface model --- frigate/data_processing/common/face/model.py | 10 ++++------ frigate/embeddings/onnx/base_embedding.py | 1 + frigate/embeddings/onnx/facenet.py | 18 +++++++++++------- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/frigate/data_processing/common/face/model.py b/frigate/data_processing/common/face/model.py index e81e0e4078..2c01142886 100644 --- a/frigate/data_processing/common/face/model.py +++ b/frigate/data_processing/common/face/model.py @@ -44,6 +44,7 @@ class FaceRecognizer(ABC): output_height: int, ) -> np.ndarray: # landmark is run on grayscale images + if image.ndim == 3: land_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: @@ -230,7 +231,6 @@ class FaceNetRecognizer(FaceRecognizer): if img is None: continue - img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img = self.align_face(img, img.shape[1], img.shape[0]) emb = self.face_embedder([img])[0].squeeze() face_embeddings_map[name].append(emb) @@ -241,8 +241,7 @@ class FaceNetRecognizer(FaceRecognizer): return for name, embs in face_embeddings_map.items(): - norms = np.linalg.norm(embs, axis=-1, keepdims=True) - self.mean_embs[name] = stats.trim_mean(embs / norms, 0.15) + self.mean_embs[name] = stats.trim_mean(embs, 0.15) def classify(self, face_image): if not self.landmark_detector: @@ -255,14 +254,13 @@ class FaceNetRecognizer(FaceRecognizer): return None # face recognition is best run on grayscale images - img = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY) # get blur factor before aligning face - blur_factor = self.get_blur_factor(img) + blur_factor = self.get_blur_factor(face_image) logger.debug(f"face detected with bluriness {blur_factor}") # align face and run recognition - img = self.align_face(img, img.shape[1], img.shape[0]) + img = self.align_face(face_image, face_image.shape[1], face_image.shape[0]) embedding = self.face_embedder([img])[0].squeeze() score = 0 diff --git a/frigate/embeddings/onnx/base_embedding.py b/frigate/embeddings/onnx/base_embedding.py index 7403f0ac1e..84f68cf429 100644 --- a/frigate/embeddings/onnx/base_embedding.py +++ b/frigate/embeddings/onnx/base_embedding.py @@ -85,6 +85,7 @@ class BaseEmbedding(ABC): input_names = self.runner.get_input_names() onnx_inputs = {name: [] for name in input_names} input: dict[str, any] + print(f"onnx inputs are {input_names}") for input in processed: for key, value in input.items(): if key in input_names: diff --git a/frigate/embeddings/onnx/facenet.py b/frigate/embeddings/onnx/facenet.py index 7041f27147..6c44b8f590 100644 --- a/frigate/embeddings/onnx/facenet.py +++ b/frigate/embeddings/onnx/facenet.py @@ -13,7 +13,7 @@ from .runner import ONNXModelRunner logger = logging.getLogger(__name__) -FACE_EMBEDDING_SIZE = 160 +FACE_EMBEDDING_SIZE = 112 class FaceNetEmbedding(BaseEmbedding): @@ -77,9 +77,9 @@ class FaceNetEmbedding(BaseEmbedding): og = np.array(pil).astype(np.float32) # Image must be FACE_EMBEDDING_SIZExFACE_EMBEDDING_SIZE - og_h, og_w = og.shape + og_h, og_w, channels = og.shape frame = np.zeros( - (FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, 3), dtype=np.float32 + (FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, channels), dtype=np.float32 ) # compute center offset @@ -87,8 +87,12 @@ class FaceNetEmbedding(BaseEmbedding): y_center = (FACE_EMBEDDING_SIZE - og_h) // 2 # copy img image into center of result image - frame[y_center : y_center + og_h, x_center : x_center + og_w, 0] = og - frame[y_center : y_center + og_h, x_center : x_center + og_w, 1] = og - frame[y_center : y_center + og_h, x_center : x_center + og_w, 2] = og + frame[y_center : y_center + og_h, x_center : x_center + og_w] = og + + # run arcface normalization + normalized_image = frame.astype(np.float32) / 255.0 + frame = (normalized_image - 0.5) / 0.5 + + frame = np.transpose(frame, (2, 0, 1)) frame = np.expand_dims(frame, axis=0) - return [{"input_2": frame}] + return [{"data": frame}]