Create face_recognition_system.py

face_recognition_system.py

@@ -0,0 +1,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()