recognition_system/src/enroll_no_gui.py
2025-05-16 11:53:34 +01:00

119 lines
3.7 KiB
Python

import cv2
import os
from datetime import datetime
import time
import argparse
import subprocess
import dlib
from imutils import face_utils
import pickle
from cryptography.fernet import Fernet
from dotenv import load_dotenv
#if "ENCODINGS_KEY" not in os.environ:
# print("Please set the ENCODINGS_KEY environment variable with your encryption key.")
# exit(1)
#key = os.environ.get("ENCODINGS_KEY").encode()
load_dotenv()
key = os.getenv("ENCODINGS_KEY")
fernet = Fernet(key)
def encrypt_data(data):
serialized = pickle.dumps(data)
encrypted = fernet.encrypt(serialized)
return encrypted
def decrypt_data(encrypted_data):
decrypted = fernet.decrypt(encrypted_data)
data = pickle.loads(decrypted)
return data
detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
def create_folder(name):
dataset_folder = "dataset"
if not os.path.exists(dataset_folder):
os.makedirs(dataset_folder)
person_folder = os.path.join(dataset_folder, name)
if not os.path.exists(person_folder):
os.makedirs(person_folder)
return person_folder
def is_blinking(landmarks):
left_eye = landmarks[36:42]
right_eye = landmarks[42:48]
def eye_aspect_ratio(eye):
A = cv2.norm(eye[1] - eye[5])
B = cv2.norm(eye[2] - eye[4])
C = cv2.norm(eye[0] - eye[3])
return (A + B) / (2.0 * C)
left_eye_ratio = eye_aspect_ratio(left_eye)
right_eye_ratio = eye_aspect_ratio(right_eye)
ear = (left_eye_ratio + right_eye_ratio) / 2.0
return ear < 0.5
def capture_photos(name, auto=False):
folder = create_folder(name)
cap = cv2.VideoCapture(2)
if not cap.isOpened():
print("Error: Could not open camera.")
return
time.sleep(2)
photo_count = 0
if not auto:
print(f"Press ENTER to start taking 10 photos for {name}. Press 'q' to quit.")
input()
else:
print(f"Starting photo capture for {name} automatically.")
while photo_count < 10:
ret, frame = cap.read()
if not ret:
print("Error: Could not read frame.")
break
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
faces = detector(gray, 0)
if len(faces) == 0:
print("No face detected. Please align properly.")
time.sleep(1)
continue
for face in faces:
landmarks = predictor(gray, face)
landmarks = face_utils.shape_to_np(landmarks)
if len(landmarks) == 68 and is_blinking(landmarks):
photo_count += 1
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"{name}_{timestamp}_{photo_count}.jpg"
filepath = os.path.join(folder, filename)
cv2.imwrite(filepath, frame)
print(f"Photo {photo_count}/10 saved: {filepath}")
else:
print("No blink detected. Try again.")
time.sleep(1)
cap.release()
print(f"Photo capture completed. {photo_count} photos saved for {name}.")
print("[INFO] Running model training script with encryption...")
venv_python = "./biometric_system/bin/python3.10"
subprocess.run([venv_python, "model_training.py"])
print("[INFO] Model training complete. Exiting...")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Capture photos for a person.")
parser.add_argument('person_name', type=str, help="The name of the person to photograph")
parser.add_argument('--auto', action="store_true", help="Run automatically without waiting for keyboard input")
args = parser.parse_args()
capture_photos(args.person_name, auto=args.auto)