Spaces:
Sleeping
Sleeping
File size: 2,704 Bytes
c77ba39 9651a0f 957d84b c77ba39 befb45d c77ba39 befb45d c77ba39 befb45d c77ba39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
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]
|