Spaces:
Sleeping
Sleeping
| import os | |
| from pathlib import Path | |
| from typing import List, Dict, Union | |
| import cv2 | |
| import numpy as np | |
| from deepface import DeepFace | |
| # Avoid RetinaFace/tf-keras mismatch: use OpenCV face detector backend | |
| DETECTOR = os.getenv("FACE_DETECTOR", "opencv") | |
| MODEL_NAME = os.getenv("FACE_MODEL", "VGG-Face") # or "Facenet512", "ArcFace", ... | |
| DIST_THRESHOLD = float(os.getenv("FACE_DIST_THRESHOLD", "0.35")) # lower => stricter | |
| def _list_images(folder: Path) -> list[Path]: | |
| exts = {".jpg", ".jpeg", ".png", ".bmp"} | |
| return [p for p in folder.glob("*") if p.suffix.lower() in exts] | |
| def _ensure_faces_dir(dir_path: str) -> Path: | |
| p = Path(dir_path) | |
| if not p.exists() or not any(p.iterdir()): | |
| # Return as-is; caller will print friendly message | |
| return p | |
| return p | |
| def _embed(img_bgr: np.ndarray): | |
| # DeepFace.represent expects RGB | |
| rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) | |
| reps = DeepFace.represent( | |
| rgb, model_name=MODEL_NAME, detector_backend=DETECTOR, enforce_detection=False | |
| ) | |
| # represent returns list of dicts. If none detected, empty list. | |
| if not reps: | |
| return None | |
| # take first face (for counting/verification, that’s fine) | |
| return np.array(reps[0]["embedding"], dtype=np.float32) | |
| def _cosine(a, b): | |
| a = a / (np.linalg.norm(a) + 1e-9) | |
| b = b / (np.linalg.norm(b) + 1e-9) | |
| return float(np.dot(a, b)) | |
| def recognize_faces(frame_bgr: np.ndarray, faces_dir: str, topk: int = 3) -> Union[str, List[Dict]]: | |
| """ | |
| Returns: | |
| - str message if faces dir missing/empty or no faces detected. | |
| - List[{"name": str, "score": float}] otherwise (score = 1 - cosine distance). | |
| """ | |
| gallery_root = _ensure_faces_dir(faces_dir) | |
| if not gallery_root.exists(): | |
| return f"Warning: faces folder '{faces_dir}' not found." | |
| gallery_imgs = _list_images(gallery_root) | |
| if not gallery_imgs: | |
| return f"Warning: faces folder '{faces_dir}' is empty." | |
| # embed incoming frame (first face) | |
| probe = _embed(frame_bgr) | |
| if probe is None: | |
| return "No face in frame" | |
| # compare to gallery | |
| scores = [] | |
| for p in gallery_imgs: | |
| img = cv2.imread(str(p)) | |
| if img is None: | |
| continue | |
| emb = _embed(img) | |
| if emb is None: | |
| continue | |
| cos = _cosine(probe, emb) | |
| # convert to distance-like: higher is better (similarity) | |
| scores.append({"name": p.stem, "score": cos}) | |
| scores.sort(key=lambda x: x["score"], reverse=True) | |
| # Filter with threshold if provided (using cosine similarity ~ 1.0 is perfect) | |
| filtered = [s for s in scores if (1.0 - s["score"]) <= DIST_THRESHOLD] | |
| return filtered[:topk] | |