diff --git a/frigate/config/classification.py b/frigate/config/classification.py index 0070569a84..25c3795465 100644 --- a/frigate/config/classification.py +++ b/frigate/config/classification.py @@ -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, diff --git a/frigate/data_processing/common/face/model.py b/frigate/data_processing/common/face/model.py new file mode 100644 index 0000000000..dc363a11c3 --- /dev/null +++ b/frigate/data_processing/common/face/model.py @@ -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) diff --git a/frigate/data_processing/real_time/face.py b/frigate/data_processing/real_time/face.py index 7b49a2f472..99c696a1aa 100644 --- a/frigate/data_processing/real_time/face.py +++ b/frigate/data_processing/real_time/face.py @@ -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 diff --git a/frigate/embeddings/onnx/base_embedding.py b/frigate/embeddings/onnx/base_embedding.py index a2ea926743..7403f0ac1e 100644 --- a/frigate/embeddings/onnx/base_embedding.py +++ b/frigate/embeddings/onnx/base_embedding.py @@ -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 diff --git a/frigate/embeddings/onnx/facenet.py b/frigate/embeddings/onnx/facenet.py new file mode 100644 index 0000000000..f32ac8d4d7 --- /dev/null +++ b/frigate/embeddings/onnx/facenet.py @@ -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}]