0
0
Fork 0
mirror of https://github.com/boltgolt/howdy.git synced 2024-09-19 09:51:19 +02:00
howdy/learn.py

85 lines
2.2 KiB
Python
Raw Normal View History

2018-01-05 16:37:00 +01:00
# Save the face of the user in encoded form
# Import required modules
2018-01-05 14:34:18 +01:00
import face_recognition
2018-01-05 01:59:44 +01:00
import subprocess
import time
import os
2018-01-05 14:34:18 +01:00
import sys
import json
2018-01-05 16:37:00 +01:00
# Import config and extra functions
2018-01-05 14:34:18 +01:00
import config
import utils
def captureFrame(delay):
2018-01-05 16:37:00 +01:00
"""Capture and encode 1 frame of video"""
# Call fswebcam to save a frame to /tmp with a set delay
2018-01-05 14:34:18 +01:00
subprocess.call(["fswebcam", "-S", str(delay), "--no-banner", "-d", "/dev/video" + str(config.device_id), tmp_file], stderr=open(os.devnull, "wb"))
2018-01-05 16:37:00 +01:00
# Get the faces in htat image
2018-01-05 14:34:18 +01:00
ref = face_recognition.load_image_file(tmp_file)
enc = face_recognition.face_encodings(ref)
2018-01-05 16:37:00 +01:00
# If 0 faces are detected we can't continue
2018-01-05 14:34:18 +01:00
if len(enc) == 0:
print("No face detected, aborting")
sys.exit()
2018-01-05 16:37:00 +01:00
# If more than 1 faces are detected we can't know wich one belongs to the user
2018-01-05 14:34:18 +01:00
if len(enc) > 1:
print("Multiple faces detected, aborting")
sys.exit()
clean_enc = []
2018-01-05 16:37:00 +01:00
# Copy the values into a clean array so we can export it as JSON later on
2018-01-05 14:34:18 +01:00
for point in enc[0]:
clean_enc.append(point)
encodings.append(clean_enc)
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# The current user
2018-01-05 01:59:44 +01:00
user = os.environ.get("USER")
2018-01-05 16:37:00 +01:00
# The name of the tmp frame file to user
2018-01-05 14:34:18 +01:00
tmp_file = "/tmp/howdy_" + user + ".jpg"
2018-01-05 16:37:00 +01:00
# The permanent file to store the encoded model in
2018-01-05 14:34:18 +01:00
enc_file = "./models/" + user + ".dat"
2018-01-05 16:37:00 +01:00
# Known encodings
2018-01-05 14:34:18 +01:00
encodings = []
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# Make the ./models folder if it doesn't already exist
2018-01-05 01:59:44 +01:00
if not os.path.exists("models"):
print("No face model folder found, creating one")
2018-01-05 14:34:18 +01:00
os.makedirs("models")
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# To try read a premade encodings file if it exists
2018-01-05 14:34:18 +01:00
try:
encodings = json.load(open(enc_file))
except FileNotFoundError:
encodings = False
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# If a file does exist, ask the user what needs to be done
2018-01-05 14:34:18 +01:00
if encodings != False:
encodings = utils.print_menu(encodings)
2018-01-05 01:59:44 +01:00
2018-01-05 14:34:18 +01:00
print("\nLearning face for the user account " + user)
print("Please look straight into the camera for 5 seconds")
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# Give the user time to read
2018-01-05 14:34:18 +01:00
time.sleep(2)
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# Capture with 3 different delays to simulate different camera exposures
2018-01-05 14:34:18 +01:00
for delay in [30, 6, 0]:
time.sleep(.3)
captureFrame(delay)
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# Save the new encodings to disk
2018-01-05 14:34:18 +01:00
with open(enc_file, "w") as datafile:
json.dump(encodings, datafile)
2018-01-05 01:59:44 +01:00
2018-01-05 16:37:00 +01:00
# Remove any left over temp files
2018-01-05 14:34:18 +01:00
os.remove(tmp_file)
2018-01-05 01:59:44 +01:00
print("Done.")