Spaces:
Sleeping
Sleeping
| 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() |