File size: 4,565 Bytes
9178542 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import cv2
import numpy as np
from ultralytics import YOLO
from arcface import ArcFace
import pickle
import os
from datetime import datetime
class FaceRecognitionSystem:
def __init__(self, database_path="face_database.pkl", confidence_threshold=0.5, similarity_threshold=2):
# Initialize YOLO for face detection
self.yolo_model = YOLO('https://github.com/akanametov/yolo-face/releases/download/v0.0.0/yolov11s-face.pt')
# Initialize ArcFace for face recognition
self.face_rec = ArcFace.ArcFace("model.tflite")
# Thresholds
self.confidence_threshold = confidence_threshold
self.similarity_threshold = similarity_threshold
# Load or create face database
self.database_path = database_path
self.face_database = self.load_database()
def load_database(self):
if os.path.exists(self.database_path):
with open(self.database_path, 'rb') as f:
return pickle.load(f)
return {}
def save_database(self):
with open(self.database_path, 'wb') as f:
pickle.dump(self.face_database, f)
def add_face_to_database(self, name, frame):
"""Add a new face to the database"""
try:
embedding = self.face_rec.calc_emb(frame)
self.face_database[name] = embedding
self.save_database()
return True
except Exception as e:
print(f"Error adding face to database: {e}")
return False
def find_closest_match(self, embedding):
"""Find the closest matching face in the database"""
if not self.face_database:
return "Unknown", 1.0
min_distance = 10000
closest_name = "Unknown"
for name, stored_embedding in self.face_database.items():
distance = self.face_rec.get_distance_embeddings(embedding, stored_embedding)
if distance < min_distance:
min_distance = distance
closest_name = name
return closest_name, min_distance
def process_frame(self, frame):
"""Process a single frame"""
# Run YOLO detection
results = self.yolo_model(frame, verbose=False)[0]
# Process each detected face
for detection in results.boxes.data:
x1, y1, x2, y2, conf, _ = detection
if conf < self.confidence_threshold:
continue
# Convert coordinates to integers
x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
# Extract face region
face_region = frame[y1:y2, x1:x2]
try:
# Calculate face embedding
embedding = self.face_rec.calc_emb(face_region)
# Find closest match
name, distance = self.find_closest_match(embedding)
# Determine if match is close enough
if distance > self.similarity_threshold:
name = "Unknown"
# Draw rectangle and name
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, f"{name} ({conf:.2f})",
(x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
0.5, (0, 255, 0), 2)
except Exception as e:
print(f"Error processing face: {e}")
return frame
def run(self):
"""Run the face recognition system on webcam feed"""
cap = cv2.VideoCapture(0)
while True:
ret, frame = cap.read()
if not ret:
break
# Process the frame
processed_frame = self.process_frame(frame)
# Display the result
cv2.imshow('Face Recognition', processed_frame)
key = cv2.waitKey(1)
if key == ord('q'):
break
elif key == ord('a'):
# Add new face to database
name = input("Enter name for new face: ")
if self.add_face_to_database(name, frame):
print(f"Successfully added {name} to database")
else:
print("Failed to add face to database")
cap.release()
cv2.destroyAllWindows() |