Spaces:
Build error
Build error
| import cv2 | |
| import os | |
| import insightface | |
| import numpy as np | |
| from collections import deque | |
| # Initialize face detector and recognizer | |
| model = insightface.app.FaceAnalysis(allowed_modules=['detection', 'recognition']) | |
| model.prepare(ctx_id=0, det_size=(640, 640)) # Use ctx_id=-1 for CPU | |
| def normalize(v): | |
| return v / np.linalg.norm(v) | |
| def cosine_similarity(a, b): | |
| return np.dot(a, b) | |
| # Load known faces | |
| known_embs = [] | |
| names = [] | |
| for fname in os.listdir("images"): | |
| if fname.lower().endswith(('.jpg', '.png')): | |
| img = cv2.imread(os.path.join("images", fname)) | |
| faces = model.get(img) | |
| if faces: | |
| emb = normalize(faces[0].embedding) | |
| known_embs.append(emb) | |
| names.append(os.path.splitext(fname)[0]) | |
| print(f"Loaded {fname}") | |
| else: | |
| print(f"No face in {fname}") | |
| # Smoothing buffers for each face | |
| face_buffers = {} # key: face_id, value: deque of embeddings | |
| # Start video | |
| cap = cv2.VideoCapture(0) | |
| face_id_counter = 0 | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| faces = model.get(frame) | |
| current_buffers = {} | |
| for i, face in enumerate(faces): | |
| x1, y1, x2, y2 = face.bbox.astype(int) | |
| emb = normalize(face.embedding) | |
| # Use a temporary ID for each face based on bbox location | |
| face_id = f"{x1}-{y1}-{x2}-{y2}" | |
| if face_id not in face_buffers: | |
| face_buffers[face_id] = deque(maxlen=5) | |
| face_buffers[face_id].append(emb) | |
| current_buffers[face_id] = face_buffers[face_id] # mark as active this frame | |
| # Smooth embedding | |
| avg_emb = normalize(np.mean(face_buffers[face_id], axis=0)) | |
| # Find best match | |
| sims = [cosine_similarity(avg_emb, known) for known in known_embs] | |
| max_idx = np.argmax(sims) | |
| name = "Unknown" | |
| if sims[max_idx] > 0.5: | |
| name = names[max_idx] | |
| # Draw result | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 200, 0), 2) | |
| cv2.putText(frame, name, (x1, y1 - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 200, 0), 2) | |
| # Remove stale buffers | |
| face_buffers = {fid: buf for fid, buf in face_buffers.items() if fid in current_buffers} | |
| # Show frame | |
| cv2.imshow("InsightFace Multi-Face Recognition", frame) | |
| if cv2.waitKey(1) == 27: # ESC key | |
| break | |
| cap.release() | |
| cv2.destroyAllWindows() | |