Face recognition bug fixes (#17401)

* Simplify normalization * Fix confidence check
2026-10-14 10:42:49 +03:00 · 2025-03-26 19:43:10 -05:00
parent a37f804469
commit ff34739f40
2 changed files with 3 additions and 4 deletions
@@ -0,0 +1,97 @@
+"""Facenet Embeddings."""
+
+import logging
+import os
+
+import numpy as np
+
+from frigate.const import MODEL_CACHE_DIR
+from frigate.util.downloader import ModelDownloader
+
+from .base_embedding import BaseEmbedding
+from .runner import ONNXModelRunner
+
+logger = logging.getLogger(__name__)
+
+FACE_EMBEDDING_SIZE = 112
+
+
+class ArcfaceEmbedding(BaseEmbedding):
+    def __init__(
+        self,
+        device: str = "AUTO",
+    ):
+        super().__init__(
+            model_name="facedet",
+            model_file="arcface.onnx",
+            download_urls={
+                "arcface.onnx": "https://github.com/NickM-27/facenet-onnx/releases/download/v1.0/arcface.onnx",
+            },
+        )
+        self.device = device
+        self.download_path = os.path.join(MODEL_CACHE_DIR, self.model_name)
+        self.tokenizer = None
+        self.feature_extractor = None
+        self.runner = None
+        files_names = list(self.download_urls.keys())
+
+        if not all(
+            os.path.exists(os.path.join(self.download_path, n)) for n in files_names
+        ):
+            logger.debug(f"starting model download for {self.model_name}")
+            self.downloader = ModelDownloader(
+                model_name=self.model_name,
+                download_path=self.download_path,
+                file_names=files_names,
+                download_func=self._download_model,
+            )
+            self.downloader.ensure_model_files()
+        else:
+            self.downloader = None
+            self._load_model_and_utils()
+            logger.debug(f"models are already downloaded for {self.model_name}")
+
+    def _load_model_and_utils(self):
+        if self.runner is None:
+            if self.downloader:
+                self.downloader.wait_for_download()
+
+            self.runner = ONNXModelRunner(
+                os.path.join(self.download_path, self.model_file),
+                self.device,
+            )
+
+    def _preprocess_inputs(self, raw_inputs):
+        pil = self._process_image(raw_inputs[0])
+
+        # handle images larger than input size
+        width, height = pil.size
+        if width != FACE_EMBEDDING_SIZE or height != FACE_EMBEDDING_SIZE:
+            if width > height:
+                new_height = int(((height / width) * FACE_EMBEDDING_SIZE) // 4 * 4)
+                pil = pil.resize((FACE_EMBEDDING_SIZE, new_height))
+            else:
+                new_width = int(((width / height) * FACE_EMBEDDING_SIZE) // 4 * 4)
+                pil = pil.resize((new_width, FACE_EMBEDDING_SIZE))
+
+        og = np.array(pil).astype(np.float32)
+
+        # Image must be FACE_EMBEDDING_SIZExFACE_EMBEDDING_SIZE
+        og_h, og_w, channels = og.shape
+        frame = np.zeros(
+            (FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, channels), dtype=np.float32
+        )
+
+        # compute center offset
+        x_center = (FACE_EMBEDDING_SIZE - og_w) // 2
+        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] = og
+
+        # run arcface normalization
+        frame = (frame / 127.5) - 1.0
+
+        frame = np.transpose(frame, (2, 0, 1))
+        frame = np.expand_dims(frame, axis=0)
+        return [{"data": frame}]