mirror of
				https://github.com/blakeblackshear/frigate.git
				synced 2025-10-27 10:52:11 +01:00 
			
		
		
		
	Face recognition improvements (#16034)
This commit is contained in:
		
							parent
							
								
									c7c8575c9b
								
							
						
					
					
						commit
						3f99ff65ed
					
				| @ -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=( | ||||
|             { | ||||
|  | ||||
| @ -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" | ||||
|  | ||||
| @ -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: | ||||
|  | ||||
| @ -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): | ||||
|  | ||||
| @ -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: | ||||
|  | ||||
| @ -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: | ||||
|  | ||||
| @ -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 | ||||
|  | ||||
		Loading…
	
		Reference in New Issue
	
	Block a user