From 0ec536a4e534fb498cfbb2f4b43a3c621c289b0f Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Sat, 18 Jan 2025 10:52:01 -0700 Subject: [PATCH] Face recognition improvements (#16034) --- frigate/api/classification.py | 14 +++- frigate/comms/embeddings_updater.py | 1 + frigate/data_processing/real_time/api.py | 5 +- .../real_time/bird_processor.py | 2 +- .../real_time/face_processor.py | 74 ++++++++++--------- frigate/embeddings/__init__.py | 5 ++ frigate/embeddings/maintainer.py | 2 +- 7 files changed, 64 insertions(+), 39 deletions(-) diff --git a/frigate/api/classification.py b/frigate/api/classification.py index 3c505d367..63f037ec2 100644 --- a/frigate/api/classification.py +++ b/frigate/api/classification.py @@ -24,14 +24,18 @@ def get_faces(): face_dict: dict[str, list[str]] = {} for name in os.listdir(FACE_DIR): - face_dict[name] = [] - face_dir = os.path.join(FACE_DIR, name) if not os.path.isdir(face_dir): continue - for file in os.listdir(face_dir): + face_dict[name] = [] + + for file in sorted( + os.listdir(face_dir), + key=lambda f: os.path.getctime(os.path.join(face_dir, f)), + reverse=True, + ): face_dict[name].append(file) return JSONResponse(status_code=200, content=face_dict) @@ -81,6 +85,10 @@ def train_face(request: Request, name: str, body: dict = None): new_name = f"{name}-{rand_id}.webp" new_file = os.path.join(FACE_DIR, f"{name}/{new_name}") shutil.move(training_file, new_file) + + context: EmbeddingsContext = request.app.embeddings + context.clear_face_classifier() + return JSONResponse( content=( { diff --git a/frigate/comms/embeddings_updater.py b/frigate/comms/embeddings_updater.py index 095f33fde..a2d0f8b29 100644 --- a/frigate/comms/embeddings_updater.py +++ b/frigate/comms/embeddings_updater.py @@ -9,6 +9,7 @@ SOCKET_REP_REQ = "ipc:///tmp/cache/embeddings" class EmbeddingsRequestEnum(Enum): + clear_face_classifier = "clear_face_classifier" embed_description = "embed_description" embed_thumbnail = "embed_thumbnail" generate_search = "generate_search" diff --git a/frigate/data_processing/real_time/api.py b/frigate/data_processing/real_time/api.py index 7f80b5287..205431a36 100644 --- a/frigate/data_processing/real_time/api.py +++ b/frigate/data_processing/real_time/api.py @@ -32,9 +32,12 @@ class RealTimeProcessorApi(ABC): pass @abstractmethod - def handle_request(self, request_data: dict[str, any]) -> dict[str, any] | None: + def handle_request( + self, topic: str, request_data: dict[str, any] + ) -> dict[str, any] | None: """Handle metadata requests. Args: + topic (str): topic that dictates what work is requested. request_data (dict): containing data about requested change to process. Returns: diff --git a/frigate/data_processing/real_time/bird_processor.py b/frigate/data_processing/real_time/bird_processor.py index e432a186b..1199f6124 100644 --- a/frigate/data_processing/real_time/bird_processor.py +++ b/frigate/data_processing/real_time/bird_processor.py @@ -146,7 +146,7 @@ class BirdProcessor(RealTimeProcessorApi): if resp.status_code == 200: self.detected_birds[obj_data["id"]] = score - def handle_request(self, request_data): + def handle_request(self, topic, request_data): return None def expire_object(self, object_id): diff --git a/frigate/data_processing/real_time/face_processor.py b/frigate/data_processing/real_time/face_processor.py index 2b12e9994..5b0d69179 100644 --- a/frigate/data_processing/real_time/face_processor.py +++ b/frigate/data_processing/real_time/face_processor.py @@ -12,6 +12,7 @@ import cv2 import numpy as np import requests +from frigate.comms.embeddings_updater import EmbeddingsRequestEnum from frigate.config import FrigateConfig from frigate.const import FACE_DIR, FRIGATE_LOCALHOST, MODEL_CACHE_DIR from frigate.util.image import area @@ -353,45 +354,52 @@ class FaceProcessor(RealTimeProcessorApi): self.__update_metrics(datetime.datetime.now().timestamp() - start) - def handle_request(self, request_data) -> dict[str, any] | None: - rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6)) - label = request_data["face_name"] - id = f"{label}-{rand_id}" - - if request_data.get("cropped"): - thumbnail = request_data["image"] - else: - img = cv2.imdecode( - np.frombuffer(base64.b64decode(request_data["image"]), dtype=np.uint8), - cv2.IMREAD_COLOR, + def handle_request(self, topic, request_data) -> dict[str, any] | None: + if topic == EmbeddingsRequestEnum.clear_face_classifier.value: + self.__clear_classifier() + elif topic == EmbeddingsRequestEnum.register_face.value: + rand_id = "".join( + random.choices(string.ascii_lowercase + string.digits, k=6) ) - face_box = self.__detect_face(img) + label = request_data["face_name"] + id = f"{label}-{rand_id}" - if not face_box: - return { - "message": "No face was detected.", - "success": False, - } + if request_data.get("cropped"): + thumbnail = request_data["image"] + else: + img = cv2.imdecode( + np.frombuffer( + base64.b64decode(request_data["image"]), dtype=np.uint8 + ), + cv2.IMREAD_COLOR, + ) + face_box = self.__detect_face(img) - face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]] - ret, thumbnail = cv2.imencode( - ".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100] - ) + if not face_box: + return { + "message": "No face was detected.", + "success": False, + } - # write face to library - folder = os.path.join(FACE_DIR, label) - file = os.path.join(folder, f"{id}.webp") - os.makedirs(folder, exist_ok=True) + face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]] + _, thumbnail = cv2.imencode( + ".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100] + ) - # save face image - with open(file, "wb") as output: - output.write(thumbnail.tobytes()) + # write face to library + folder = os.path.join(FACE_DIR, label) + file = os.path.join(folder, f"{id}.webp") + os.makedirs(folder, exist_ok=True) - self.__clear_classifier() - return { - "message": "Successfully registered face.", - "success": True, - } + # save face image + with open(file, "wb") as output: + output.write(thumbnail.tobytes()) + + self.__clear_classifier() + return { + "message": "Successfully registered face.", + "success": True, + } def expire_object(self, object_id: str): if object_id in self.detected_faces: diff --git a/frigate/embeddings/__init__.py b/frigate/embeddings/__init__.py index dd05fb0ca..4a3f898e7 100644 --- a/frigate/embeddings/__init__.py +++ b/frigate/embeddings/__init__.py @@ -211,6 +211,11 @@ class EmbeddingsContext: return self.db.execute_sql(sql_query).fetchall() + def clear_face_classifier(self) -> None: + self.requestor.send_data( + EmbeddingsRequestEnum.clear_face_classifier.value, None + ) + def delete_face_ids(self, face: str, ids: list[str]) -> None: folder = os.path.join(FACE_DIR, face) for id in ids: diff --git a/frigate/embeddings/maintainer.py b/frigate/embeddings/maintainer.py index aa0322fd7..e221bd146 100644 --- a/frigate/embeddings/maintainer.py +++ b/frigate/embeddings/maintainer.py @@ -140,7 +140,7 @@ class EmbeddingMaintainer(threading.Thread): ) else: for processor in self.processors: - resp = processor.handle_request(data) + resp = processor.handle_request(topic, data) if resp is not None: return resp