| import pickle |
| from pathlib import Path |
| from typing import List, Optional |
|
|
| import cv2 |
| import face_recognition |
| import numpy as np |
| from PIL import Image |
| from pydantic import model_validator |
|
|
| from ....models.od.schemas import Target |
| from ....utils.registry import registry |
| from ...base import ArgSchema, BaseModelTool |
|
|
| ARGSCHEMA = {} |
|
|
|
|
| @registry.register_tool() |
| class FaceRecognition(BaseModelTool): |
| args_schema: ArgSchema = ArgSchema(**ARGSCHEMA) |
| description: str = ( |
| "This tool can recognize facial information in images/extracted frames and identify who the person is." |
| "Images/extracted frames are already loaded." |
| ) |
| threshold: float = 0.6 |
| num_jitters: int = 1 |
| face_db: str = "data/face_db" |
| model: str = "large" |
| loaded_face_db: Optional[dict] = None |
|
|
| @model_validator(mode="after") |
| def face_db_validator(self) -> "FaceRecognition": |
| if self.loaded_face_db is None: |
| if Path(self.face_db).exists(): |
| self.loaded_face_db = self._load_face_db(self.face_db) |
| else: |
| raise ValueError(f"Face database not found at {self.face_db}") |
| elif isinstance(self.loaded_face_db, dict): |
| if ( |
| "embeddings" not in self.loaded_face_db |
| or "names" not in self.loaded_face_db |
| ): |
| raise ValueError( |
| "Face database must have 'embeddings' and 'names' keys." |
| ) |
| else: |
| raise ValueError("Face database must be a dictionary.") |
| return self |
|
|
| def _load_face_db(self, path: str): |
| cached_model = Path(path).joinpath(f"representations_{self.model}_face.pkl") |
| |
| |
| |
| face_db = Path(path) |
| embeddings = [] |
| names = [] |
| for known_image in face_db.rglob("*"): |
| if known_image.suffix in [".jpg", ".png", ".webp"]: |
| loaded_image = np.array(Image.open(known_image).convert("RGB")) |
| loaded_image = cv2.cvtColor(loaded_image, cv2.COLOR_RGB2BGR) |
| known_encoding = face_recognition.face_encodings( |
| loaded_image, model="large" |
| )[0] |
| embeddings.append(known_encoding) |
| names.append(known_image.parent.name) |
| loaded_face_db = {"embeddings": embeddings, "names": names} |
| pickle.dump(loaded_face_db, open(cached_model, "wb")) |
| return loaded_face_db |
|
|
| def infer(self, image: Image.Image) -> List[Target]: |
| img = np.array(image.convert("RGB")) |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) |
| face_locations = face_recognition.face_locations(img) |
| face_encodings = face_recognition.face_encodings(img, face_locations) |
| rec_res = [] |
| for (top, right, bottom, left), face_encoding in zip( |
| face_locations, face_encodings |
| ): |
| face_distances = face_recognition.face_distance( |
| self.loaded_face_db.get("embeddings"), face_encoding |
| ) |
| best_match_index = np.argmin(face_distances) |
| if face_distances[best_match_index] <= self.threshold: |
| name = self.loaded_face_db["names"][best_match_index] |
| bbox = [left, top, right, bottom] |
| rec_res.append( |
| Target(label=name, bbox=bbox, conf=face_distances[best_match_index]) |
| ) |
| return rec_res |
|
|
| def _run(self): |
| names = set() |
| for key in self.stm.image_cache.keys(): |
| anno = self.infer(self.stm.image_cache[key]) |
| self.stm.image_cache[key] = self.visual_prompting( |
| self.stm.image_cache[key], anno |
| ) |
| names.update([item.label for item in anno]) |
|
|
| return f"Recognized {len(names)} faces: {', '.join(names)}" |
|
|
| async def _arun(self): |
| return self._run() |
|
|
