From d8f897a172b2ccc7d4992d777f1dcbb06b5c36c6 Mon Sep 17 00:00:00 2001 From: Kyle Anderson Date: Mon, 28 Oct 2019 16:26:26 -0400 Subject: [PATCH] Prepare for user registration Still preparing for user registration stuff. Made some changes to authentication to make some of its general methods available for the user registration module as well. --- authenticate_user.py | 34 +++++++++------------------------- common.py | 19 +++++++++++++++++++ create_known_encoding.py | 3 --- data_handler.py | 27 +++++++++++++++++++++++++++ register_user.py | 34 ++++++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 28 deletions(-) create mode 100644 common.py delete mode 100644 create_known_encoding.py create mode 100644 data_handler.py create mode 100644 register_user.py diff --git a/authenticate_user.py b/authenticate_user.py index e7fd100..957b998 100644 --- a/authenticate_user.py +++ b/authenticate_user.py @@ -5,31 +5,22 @@ Methods for authenticating a user. from imutils.video import VideoStream import face_recognition import imutils -import pickle import time import cv2 +import common +from data_handler import load_database + # How long to wait before timing out and saying failed authentication. TIMEOUT: float = 30.0 # Minimum number of frames in which a user must be recognized in order to be authenticated. MIN_USER_RECOGNITION_COUNT = 10 -image_writer = None -USER_IDS_KEY: str = "names" # TODO change +USER_IDS_KEY: str = "user_ids" def load_encodings(file_location: str): """Loads the encodings for faces from the given file location.""" - with open(file_location, "rb") as encodings_file: - encodings = pickle.loads(encodings_file.read()) - return encodings - - -def start_video_stream(camera: int): - """Starts the video stream and returns the created stream. - Also waits for the video stream to open before returning it.""" - video_stream = VideoStream(src=camera).start() - time.sleep(2.0) - return video_stream + return load_database(file_location) def determine_identity(face_encoding, known_faces): @@ -77,15 +68,10 @@ def draw_rectanges_and_user_ids(image_frame, conversion: float, boxes, user_ids: # Find the top so we can put the text there. y = top - 15 if top - 15 > 15 else top + 15 cv2.putText(image_frame, user_id, (left, y), cv2.FONT_HERSHEY_PLAIN, 0.75, (0, 255, 0), 2) - display_frame(image_frame) - - -def display_frame(frame): - """Displays the frame to the user.""" - cv2.imshow("Frame", frame) + common.display_frame(image_frame) -def recognize_user(encodings_location: str = "./encodings.pickle", encoding_model: str = "hog", image_flip: int = None, +def recognize_user(known_faces: dict, encoding_model: str = "hog", image_flip: int = None, draw_rectangles=False): """Attempts to recognize a user. Returns the ID of the user if identified, or None if no users are identified. @@ -93,9 +79,7 @@ def recognize_user(encodings_location: str = "./encodings.pickle", encoding_mode track of how many times each user was recognized.""" recognized_users_count = {} recognized_user = None - # video_stream = start_video_stream(0) # TODO add back - video_stream = VideoStream(src=0).start() # TODO remove - known_faces = load_encodings(encodings_location) + video_stream = common.start_video_stream(0) # Determine the time at which we will time out. Equal to current time + timeout. timeout_time: float = time.time() + TIMEOUT @@ -158,7 +142,7 @@ if __name__ == "__main__": parser.add_argument("--show", "-s", action="store_true", help="Include this argument to have the image shown to you.", default=False) args = parser.parse_args() - user = recognize_user(encoding_model=args.model, encodings_location=args.encodings, image_flip=args.flip, + user = recognize_user(encoding_model=args.model, encodings=load_encodings(args.encodings), image_flip=args.flip, draw_rectangles=args.show) if user: print(f"Recognized user {user}.") diff --git a/common.py b/common.py new file mode 100644 index 0000000..e39c478 --- /dev/null +++ b/common.py @@ -0,0 +1,19 @@ +import time + +import cv2 +from imutils.video import VideoStream + +USER_IDS_KEY = "user_ids" + + +def display_frame(frame): + """Displays the frame to the user.""" + cv2.imshow("Frame", frame) + + +def start_video_stream(camera: int): + """Starts the video stream and returns the created stream. + Also waits for the video stream to open before returning it.""" + video_stream = VideoStream(src=camera).start() + time.sleep(2.0) + return video_stream diff --git a/create_known_encoding.py b/create_known_encoding.py deleted file mode 100644 index 7f24361..0000000 --- a/create_known_encoding.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Creates a facial recognition profile for a new user. -""" diff --git a/data_handler.py b/data_handler.py new file mode 100644 index 0000000..5ef06bb --- /dev/null +++ b/data_handler.py @@ -0,0 +1,27 @@ +""" +General IO for pickle database operations. +""" + +import pickle + + +def load_database(file_location: str): + """ + Attempts to load the pickle database at the given file location + :param file_location: String location of file to be loaded. + :return: The loaded pickle database. + """ + with open(file_location, "rb") as database: + file_content = pickle.load(database) + return file_content + + +def write_database(output_file: str, database_content: dict) -> None: + """ + Writes the dictionary database to the given file location + :param output_file: The location of the file to be outputted on. + :param database_content: The database content to be written to the file. + :return: None + """ + with open(output_file, "wb") as output: + output.write(pickle.dump(database_content, output)) diff --git a/register_user.py b/register_user.py new file mode 100644 index 0000000..53742c1 --- /dev/null +++ b/register_user.py @@ -0,0 +1,34 @@ +""" +Creates a facial recognition profile for a new user. +""" +from common import USER_IDS_KEY, start_video_stream +import face_recognition +import cv2 + + +def process_image(image, encoding_model: str = "hog"): + """ + Processes a single image, returning the encoded face object + :param image: The image containing the face to be processed. + :param encoding_model: The encoding model, either CNN or HOG + :return: The processed facial recognition profile encoding. + """ + image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert between the image formats + + # Detect the coordinates of the boxes corresponding to faces in the input image + boxes = face_recognition.face_locations(image_rgb, model=encoding_model) + # Actually make the encodings for the face. + return face_recognition.face_encodings(image_rgb, boxes) + + +def register_user(user_id: str, encoding_model="hog", image_flip: int = None, video_source=None): + """ + Function for registering a new user using the given video source. If video source isn't provided, then the camera + on id 0 is used. + :param user_id: The user id for the user that is being registered. + :param encoding_model: The type of encoding model. Must be either "hog" or "cnn". HOG is faster, CNN is more thorough. + :param image_flip: The integer by which this image should be flipped. 0 for reflect about x-axis, 1 for reflect on y-axis + :return: Encoded face that was detected, or None if no face was detected or if there was another error. + """ + if video_source is None: + video_source = start_video_stream(0) -- GitLab