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
							
								
									9a2de78fc9
								
							
						
					
					
						commit
						0ec536a4e5
					
				| @ -24,14 +24,18 @@ def get_faces(): | |||||||
|     face_dict: dict[str, list[str]] = {} |     face_dict: dict[str, list[str]] = {} | ||||||
| 
 | 
 | ||||||
|     for name in os.listdir(FACE_DIR): |     for name in os.listdir(FACE_DIR): | ||||||
|         face_dict[name] = [] |  | ||||||
| 
 |  | ||||||
|         face_dir = os.path.join(FACE_DIR, name) |         face_dir = os.path.join(FACE_DIR, name) | ||||||
| 
 | 
 | ||||||
|         if not os.path.isdir(face_dir): |         if not os.path.isdir(face_dir): | ||||||
|             continue |             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) |             face_dict[name].append(file) | ||||||
| 
 | 
 | ||||||
|     return JSONResponse(status_code=200, content=face_dict) |     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_name = f"{name}-{rand_id}.webp" | ||||||
|     new_file = os.path.join(FACE_DIR, f"{name}/{new_name}") |     new_file = os.path.join(FACE_DIR, f"{name}/{new_name}") | ||||||
|     shutil.move(training_file, new_file) |     shutil.move(training_file, new_file) | ||||||
|  | 
 | ||||||
|  |     context: EmbeddingsContext = request.app.embeddings | ||||||
|  |     context.clear_face_classifier() | ||||||
|  | 
 | ||||||
|     return JSONResponse( |     return JSONResponse( | ||||||
|         content=( |         content=( | ||||||
|             { |             { | ||||||
|  | |||||||
| @ -9,6 +9,7 @@ SOCKET_REP_REQ = "ipc:///tmp/cache/embeddings" | |||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| class EmbeddingsRequestEnum(Enum): | class EmbeddingsRequestEnum(Enum): | ||||||
|  |     clear_face_classifier = "clear_face_classifier" | ||||||
|     embed_description = "embed_description" |     embed_description = "embed_description" | ||||||
|     embed_thumbnail = "embed_thumbnail" |     embed_thumbnail = "embed_thumbnail" | ||||||
|     generate_search = "generate_search" |     generate_search = "generate_search" | ||||||
|  | |||||||
| @ -32,9 +32,12 @@ class RealTimeProcessorApi(ABC): | |||||||
|         pass |         pass | ||||||
| 
 | 
 | ||||||
|     @abstractmethod |     @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. |         """Handle metadata requests. | ||||||
|         Args: |         Args: | ||||||
|  |             topic (str): topic that dictates what work is requested. | ||||||
|             request_data (dict): containing data about requested change to process. |             request_data (dict): containing data about requested change to process. | ||||||
| 
 | 
 | ||||||
|         Returns: |         Returns: | ||||||
|  | |||||||
| @ -146,7 +146,7 @@ class BirdProcessor(RealTimeProcessorApi): | |||||||
|         if resp.status_code == 200: |         if resp.status_code == 200: | ||||||
|             self.detected_birds[obj_data["id"]] = score |             self.detected_birds[obj_data["id"]] = score | ||||||
| 
 | 
 | ||||||
|     def handle_request(self, request_data): |     def handle_request(self, topic, request_data): | ||||||
|         return None |         return None | ||||||
| 
 | 
 | ||||||
|     def expire_object(self, object_id): |     def expire_object(self, object_id): | ||||||
|  | |||||||
| @ -12,6 +12,7 @@ import cv2 | |||||||
| import numpy as np | import numpy as np | ||||||
| import requests | import requests | ||||||
| 
 | 
 | ||||||
|  | from frigate.comms.embeddings_updater import EmbeddingsRequestEnum | ||||||
| from frigate.config import FrigateConfig | from frigate.config import FrigateConfig | ||||||
| from frigate.const import FACE_DIR, FRIGATE_LOCALHOST, MODEL_CACHE_DIR | from frigate.const import FACE_DIR, FRIGATE_LOCALHOST, MODEL_CACHE_DIR | ||||||
| from frigate.util.image import area | from frigate.util.image import area | ||||||
| @ -353,45 +354,52 @@ class FaceProcessor(RealTimeProcessorApi): | |||||||
| 
 | 
 | ||||||
|         self.__update_metrics(datetime.datetime.now().timestamp() - start) |         self.__update_metrics(datetime.datetime.now().timestamp() - start) | ||||||
| 
 | 
 | ||||||
|     def handle_request(self, request_data) -> dict[str, any] | None: |     def handle_request(self, topic, request_data) -> dict[str, any] | None: | ||||||
|         rand_id = "".join(random.choices(string.ascii_lowercase + string.digits, k=6)) |         if topic == EmbeddingsRequestEnum.clear_face_classifier.value: | ||||||
|         label = request_data["face_name"] |             self.__clear_classifier() | ||||||
|         id = f"{label}-{rand_id}" |         elif topic == EmbeddingsRequestEnum.register_face.value: | ||||||
| 
 |             rand_id = "".join( | ||||||
|         if request_data.get("cropped"): |                 random.choices(string.ascii_lowercase + string.digits, k=6) | ||||||
|             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) |             label = request_data["face_name"] | ||||||
|  |             id = f"{label}-{rand_id}" | ||||||
| 
 | 
 | ||||||
|             if not face_box: |             if request_data.get("cropped"): | ||||||
|                 return { |                 thumbnail = request_data["image"] | ||||||
|                     "message": "No face was detected.", |             else: | ||||||
|                     "success": False, |                 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]] |                 if not face_box: | ||||||
|             ret, thumbnail = cv2.imencode( |                     return { | ||||||
|                 ".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100] |                         "message": "No face was detected.", | ||||||
|             ) |                         "success": False, | ||||||
|  |                     } | ||||||
| 
 | 
 | ||||||
|         # write face to library |                 face = img[face_box[1] : face_box[3], face_box[0] : face_box[2]] | ||||||
|         folder = os.path.join(FACE_DIR, label) |                 _, thumbnail = cv2.imencode( | ||||||
|         file = os.path.join(folder, f"{id}.webp") |                     ".webp", face, [int(cv2.IMWRITE_WEBP_QUALITY), 100] | ||||||
|         os.makedirs(folder, exist_ok=True) |                 ) | ||||||
| 
 | 
 | ||||||
|         # save face image |             # write face to library | ||||||
|         with open(file, "wb") as output: |             folder = os.path.join(FACE_DIR, label) | ||||||
|             output.write(thumbnail.tobytes()) |             file = os.path.join(folder, f"{id}.webp") | ||||||
|  |             os.makedirs(folder, exist_ok=True) | ||||||
| 
 | 
 | ||||||
|         self.__clear_classifier() |             # save face image | ||||||
|         return { |             with open(file, "wb") as output: | ||||||
|             "message": "Successfully registered face.", |                 output.write(thumbnail.tobytes()) | ||||||
|             "success": True, | 
 | ||||||
|         } |             self.__clear_classifier() | ||||||
|  |             return { | ||||||
|  |                 "message": "Successfully registered face.", | ||||||
|  |                 "success": True, | ||||||
|  |             } | ||||||
| 
 | 
 | ||||||
|     def expire_object(self, object_id: str): |     def expire_object(self, object_id: str): | ||||||
|         if object_id in self.detected_faces: |         if object_id in self.detected_faces: | ||||||
|  | |||||||
| @ -211,6 +211,11 @@ class EmbeddingsContext: | |||||||
| 
 | 
 | ||||||
|         return self.db.execute_sql(sql_query).fetchall() |         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: |     def delete_face_ids(self, face: str, ids: list[str]) -> None: | ||||||
|         folder = os.path.join(FACE_DIR, face) |         folder = os.path.join(FACE_DIR, face) | ||||||
|         for id in ids: |         for id in ids: | ||||||
|  | |||||||
| @ -140,7 +140,7 @@ class EmbeddingMaintainer(threading.Thread): | |||||||
|                     ) |                     ) | ||||||
|                 else: |                 else: | ||||||
|                     for processor in self.processors: |                     for processor in self.processors: | ||||||
|                         resp = processor.handle_request(data) |                         resp = processor.handle_request(topic, data) | ||||||
| 
 | 
 | ||||||
|                         if resp is not None: |                         if resp is not None: | ||||||
|                             return resp |                             return resp | ||||||
|  | |||||||
		Loading…
	
		Reference in New Issue
	
	Block a user