Refactor face recognition to allow for running lbph or embedding

This commit is contained in:
Nicolas Mowen 2025-03-25 09:43:43 -06:00
parent 2c3ea5b74e
commit a13494c8df
5 changed files with 392 additions and 153 deletions

View File

@ -51,6 +51,9 @@ class SemanticSearchConfig(FrigateBaseModel):
class FaceRecognitionConfig(FrigateBaseModel):
enabled: bool = Field(default=False, title="Enable face recognition.")
model_size: str = Field(
default="small", title="The size of the embeddings model used."
)
min_score: float = Field(
title="Minimum face distance score required to save the attempt.",
default=0.8,

View File

@ -0,0 +1,277 @@
import logging
import os
from abc import ABC, abstractmethod
import cv2
import numpy as np
from frigate.config import FrigateConfig
from frigate.const import MODEL_CACHE_DIR
from frigate.embeddings.onnx.facenet import FaceNetEmbedding
logger = logging.getLogger(__name__)
class FaceRecognizer(ABC):
"""Face recognition runner."""
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")
)
@abstractmethod
def build(self) -> None:
"""Build face recognition model."""
pass
@abstractmethod
def clear(self) -> None:
"""Clear current built model."""
pass
@abstractmethod
def classify(self, face_image: np.ndarray) -> tuple[str, float] | None:
pass
def align_face(
self,
image: np.ndarray,
output_width: int,
output_height: int,
) -> np.ndarray:
_, lands = self.landmark_detector.fit(
image, np.array([(0, 0, image.shape[1], image.shape[0])])
)
landmarks: np.ndarray = lands[0][0]
# get landmarks for eyes
leftEyePts = landmarks[42:48]
rightEyePts = landmarks[36:42]
# compute the center of mass for each eye
leftEyeCenter = leftEyePts.mean(axis=0).astype("int")
rightEyeCenter = rightEyePts.mean(axis=0).astype("int")
# compute the angle between the eye centroids
dY = rightEyeCenter[1] - leftEyeCenter[1]
dX = rightEyeCenter[0] - leftEyeCenter[0]
angle = np.degrees(np.arctan2(dY, dX)) - 180
# compute the desired right eye x-coordinate based on the
# desired x-coordinate of the left eye
desiredRightEyeX = 1.0 - 0.35
# determine the scale of the new resulting image by taking
# the ratio of the distance between eyes in the *current*
# image to the ratio of distance between eyes in the
# *desired* image
dist = np.sqrt((dX**2) + (dY**2))
desiredDist = desiredRightEyeX - 0.35
desiredDist *= output_width
scale = desiredDist / dist
# compute center (x, y)-coordinates (i.e., the median point)
# between the two eyes in the input image
# grab the rotation matrix for rotating and scaling the face
eyesCenter = (
int((leftEyeCenter[0] + rightEyeCenter[0]) // 2),
int((leftEyeCenter[1] + rightEyeCenter[1]) // 2),
)
M = cv2.getRotationMatrix2D(eyesCenter, angle, scale)
# update the translation component of the matrix
tX = output_width * 0.5
tY = output_height * 0.35
M[0, 2] += tX - eyesCenter[0]
M[1, 2] += tY - eyesCenter[1]
# apply the affine transformation
return cv2.warpAffine(
image, M, (output_width, output_height), flags=cv2.INTER_CUBIC
)
def get_blur_factor(self, input: np.ndarray) -> float:
"""Calculates the factor for the confidence based on the blur of the image."""
if not self.config.face_recognition.blur_confidence_filter:
return 1.0
variance = cv2.Laplacian(input, cv2.CV_64F).var()
if variance < 60: # image is very blurry
return 0.96
elif variance < 70: # image moderately blurry
return 0.98
elif variance < 80: # image is slightly blurry
return 0.99
else:
return 1.0
class LBPHRecognizer(FaceRecognizer):
def __init__(self, config: FrigateConfig):
super().__init__(config)
self.label_map: dict[int, str] = {}
self.recognizer: cv2.face.LBPHFaceRecognizer | None = None
def clear(self) -> None:
self.face_recognizer = None
self.label_map = {}
def build(self):
if not self.landmark_detector:
return None
labels = []
faces = []
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
self.label_map[idx] = name
for image in os.listdir(face_folder):
img = cv2.imread(os.path.join(face_folder, image))
if img is None:
continue
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = self.align_face(img, img.shape[1], img.shape[0])
faces.append(img)
labels.append(idx)
idx += 1
if not faces:
return
self.recognizer: cv2.face.LBPHFaceRecognizer = (
cv2.face.LBPHFaceRecognizer_create(
radius=2, threshold=(1 - self.config.face_recognition.min_score) * 1000
)
)
self.recognizer.train(faces, np.array(labels))
def classify(self, face_image: np.ndarray) -> tuple[str, float] | None:
if not self.landmark_detector:
return None
if not self.label_map or not self.recognizer:
self.build()
if not self.recognizer:
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)
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])
index, distance = self.recognizer.predict(img)
if index == -1:
return None
score = (1.0 - (distance / 1000)) * blur_factor
return self.label_map[index], round(score, 2)
class FaceNetRecognizer(FaceRecognizer):
def __init__(self, config: FrigateConfig):
super().__init__(config)
self.mean_embs: dict[int, np.ndarray] = {}
self.face_embedder: FaceNetEmbedding = FaceNetEmbedding()
def clear(self) -> None:
self.mean_embs = None
def build(self):
if not self.landmark_detector:
return None
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:
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)
idx += 1
if not face_embeddings_map:
return
for name, embs in face_embeddings_map.items():
self.mean_embs[name] = np.mean(embs, axis=0)
def classify(self, face_image):
if not self.landmark_detector:
return None
if not self.mean_embs:
self.build()
if not self.mean_embs:
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)
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])
embedding = self.face_embedder([img])[0].squeeze()
score = 0
label = ""
for name, mean_emb in self.mean_embs.items():
dot_product = np.dot(embedding, mean_emb)
magnitude_A = np.linalg.norm(embedding)
magnitude_B = np.linalg.norm(mean_emb)
cosine_similarity = dot_product / (magnitude_A * magnitude_B)
if cosine_similarity > score:
score = cosine_similarity
label = name
if score < self.config.face_recognition.min_score:
return None
return label, round(score * blur_factor, 2)

View File

@ -19,6 +19,11 @@ from frigate.comms.event_metadata_updater import (
)
from frigate.config import FrigateConfig
from frigate.const import FACE_DIR, MODEL_CACHE_DIR
from frigate.data_processing.common.face.model import (
FaceNetRecognizer,
FaceRecognizer,
LBPHRecognizer,
)
from frigate.util.image import area
from ..types import DataProcessorMetrics
@ -42,10 +47,9 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
self.face_config = config.face_recognition
self.sub_label_publisher = sub_label_publisher
self.face_detector: cv2.FaceDetectorYN = None
self.landmark_detector: cv2.face.FacemarkLBF = None
self.recognizer: cv2.face.LBPHFaceRecognizer = None
self.requires_face_detection = "face" not in self.config.objects.all_objects
self.detected_faces: dict[str, float] = {}
self.recognizer: FaceRecognizer | None = None
download_path = os.path.join(MODEL_CACHE_DIR, "facedet")
self.model_files = {
@ -72,7 +76,13 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
self.__build_detector()
self.label_map: dict[int, str] = {}
self.__build_classifier()
if self.face_config.model_size == "smal":
self.recognizer = LBPHRecognizer(self.config)
else:
self.recognizer = FaceNetRecognizer(self.config)
self.recognizer.build()
def __download_models(self, path: str) -> None:
try:
@ -92,126 +102,6 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
score_threshold=0.5,
nms_threshold=0.3,
)
self.landmark_detector = cv2.face.createFacemarkLBF()
self.landmark_detector.loadModel(
os.path.join(MODEL_CACHE_DIR, "facedet/landmarkdet.yaml")
)
def __build_classifier(self) -> None:
if not self.landmark_detector:
return None
labels = []
faces = []
dir = "/media/frigate/clips/faces"
for idx, name in enumerate(os.listdir(dir)):
if name == "train":
continue
face_folder = os.path.join(dir, name)
if not os.path.isdir(face_folder):
continue
self.label_map[idx] = name
for image in os.listdir(face_folder):
img = cv2.imread(os.path.join(face_folder, image))
if img is None:
continue
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = self.__align_face(img, img.shape[1], img.shape[0])
faces.append(img)
labels.append(idx)
if not faces:
return
self.recognizer: cv2.face.LBPHFaceRecognizer = (
cv2.face.LBPHFaceRecognizer_create(
radius=2, threshold=(1 - self.face_config.min_score) * 1000
)
)
self.recognizer.train(faces, np.array(labels))
def __align_face(
self,
image: np.ndarray,
output_width: int,
output_height: int,
) -> np.ndarray:
_, lands = self.landmark_detector.fit(
image, np.array([(0, 0, image.shape[1], image.shape[0])])
)
landmarks: np.ndarray = lands[0][0]
# get landmarks for eyes
leftEyePts = landmarks[42:48]
rightEyePts = landmarks[36:42]
# compute the center of mass for each eye
leftEyeCenter = leftEyePts.mean(axis=0).astype("int")
rightEyeCenter = rightEyePts.mean(axis=0).astype("int")
# compute the angle between the eye centroids
dY = rightEyeCenter[1] - leftEyeCenter[1]
dX = rightEyeCenter[0] - leftEyeCenter[0]
angle = np.degrees(np.arctan2(dY, dX)) - 180
# compute the desired right eye x-coordinate based on the
# desired x-coordinate of the left eye
desiredRightEyeX = 1.0 - 0.35
# determine the scale of the new resulting image by taking
# the ratio of the distance between eyes in the *current*
# image to the ratio of distance between eyes in the
# *desired* image
dist = np.sqrt((dX**2) + (dY**2))
desiredDist = desiredRightEyeX - 0.35
desiredDist *= output_width
scale = desiredDist / dist
# compute center (x, y)-coordinates (i.e., the median point)
# between the two eyes in the input image
# grab the rotation matrix for rotating and scaling the face
eyesCenter = (
int((leftEyeCenter[0] + rightEyeCenter[0]) // 2),
int((leftEyeCenter[1] + rightEyeCenter[1]) // 2),
)
M = cv2.getRotationMatrix2D(eyesCenter, angle, scale)
# update the translation component of the matrix
tX = output_width * 0.5
tY = output_height * 0.35
M[0, 2] += tX - eyesCenter[0]
M[1, 2] += tY - eyesCenter[1]
# apply the affine transformation
return cv2.warpAffine(
image, M, (output_width, output_height), flags=cv2.INTER_CUBIC
)
def __get_blur_factor(self, input: np.ndarray) -> float:
"""Calculates the factor for the confidence based on the blur of the image."""
if not self.face_config.blur_confidence_filter:
return 1.0
variance = cv2.Laplacian(input, cv2.CV_64F).var()
if variance < 60: # image is very blurry
return 0.96
elif variance < 70: # image moderately blurry
return 0.98
elif variance < 80: # image is slightly blurry
return 0.99
else:
return 1.0
def __clear_classifier(self) -> None:
self.face_recognizer = None
self.label_map = {}
def __detect_face(
self, input: np.ndarray, threshold: float
@ -254,33 +144,6 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
return face
def __classify_face(self, face_image: np.ndarray) -> tuple[str, float] | None:
if not self.landmark_detector:
return None
if not self.label_map or not self.recognizer:
self.__build_classifier()
if not self.recognizer:
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)
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])
index, distance = self.recognizer.predict(img)
if index == -1:
return None
score = (1.0 - (distance / 1000)) * blur_factor
return self.label_map[index], round(score, 2)
def __update_metrics(self, duration: float) -> None:
self.metrics.face_rec_fps.value = (
self.metrics.face_rec_fps.value * 9 + duration
@ -370,7 +233,7 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
max(0, face_box[0]) : min(frame.shape[1], face_box[2]),
]
res = self.__classify_face(face_frame)
res = self.recognizer.classify(face_frame)
if not res:
return
@ -431,7 +294,7 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
return {"message": "No face was detected.", "success": False}
face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]]
res = self.__classify_face(face)
res = self.recognizer.classify(face)
if not res:
return {"success": False, "message": "No face was recognized."}
@ -500,7 +363,7 @@ class FaceRealTimeProcessor(RealTimeProcessorApi):
"success": False,
}
res = self.__classify_face(img)
res = self.recognizer.classify(img)
if not res:
return

View File

@ -69,6 +69,8 @@ class BaseEmbedding(ABC):
image = Image.open(BytesIO(response.content)).convert(output)
elif isinstance(image, bytes):
image = Image.open(BytesIO(image)).convert(output)
elif isinstance(image, np.ndarray):
image = Image.fromarray(image)
return image

View File

@ -0,0 +1,94 @@
"""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 = 160
class FaceNetEmbedding(BaseEmbedding):
def __init__(
self,
device: str = "AUTO",
):
super().__init__(
model_name="facedet",
model_file="facenet.onnx",
download_urls={
"facenet.onnx": "https://huggingface.co/jinaai/jina-clip-v1/resolve/main/onnx/text_model_fp16.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 = og.shape
frame = np.zeros(
(FACE_EMBEDDING_SIZE, FACE_EMBEDDING_SIZE, 3), 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, 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 = np.expand_dims(frame, axis=0)
return [{"input_2": frame}]