119 lines
3.7 KiB
Python
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)
|