mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-06-22 01:18:21 +02:00
b18d1fb970
* Refactor face recognition to allow for running lbph or embedding * Cleanup * Use weighted average for faces * Set correct url * Cleanup * Update docs * Update docs * Use scipy trimmed mean * Normalize * Handle color and gray landmark detection * Upgrade to new arcface model * Implement sigmoid function * Rename * Rename to arcface * Fix * Add face recognition model size to ui config * Update toast
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
"""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
|
|
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 [{"data": frame}]
|