Commit 7533d929 authored by Kyle Anderson's avatar Kyle Anderson

Work on register user

The register user functionality isn't quite working, doesn't save to the database properly. Will continue working on this.
parent d8f897a1
...@@ -116,7 +116,7 @@ def recognize_user(known_faces: dict, encoding_model: str = "hog", image_flip: i ...@@ -116,7 +116,7 @@ def recognize_user(known_faces: dict, encoding_model: str = "hog", image_flip: i
recognized_users = check_recognized_users(recognized_users_count) recognized_users = check_recognized_users(recognized_users_count)
if len(recognized_users) > 0: if len(recognized_users) > 0:
break break
cv2.waitKey(1) # Required or else video stream doesn't really render. cv2.waitKey(1) # Required or else video stream doesn't really render.
if recognized_users_count: if recognized_users_count:
recognized_user = max(recognized_users_count, recognized_user = max(recognized_users_count,
...@@ -142,7 +142,7 @@ if __name__ == "__main__": ...@@ -142,7 +142,7 @@ if __name__ == "__main__":
parser.add_argument("--show", "-s", action="store_true", parser.add_argument("--show", "-s", action="store_true",
help="Include this argument to have the image shown to you.", default=False) help="Include this argument to have the image shown to you.", default=False)
args = parser.parse_args() args = parser.parse_args()
user = recognize_user(encoding_model=args.model, encodings=load_encodings(args.encodings), image_flip=args.flip, user = recognize_user(known_faces=load_encodings(args.encodings), encoding_model=args.model, image_flip=args.flip,
draw_rectangles=args.show) draw_rectangles=args.show)
if user: if user:
print(f"Recognized user {user}.") print(f"Recognized user {user}.")
...@@ -4,6 +4,7 @@ import cv2 ...@@ -4,6 +4,7 @@ import cv2
from imutils.video import VideoStream from imutils.video import VideoStream
USER_IDS_KEY = "user_ids" USER_IDS_KEY = "user_ids"
DATABASE_LOC = "./dataset/faces.pickle"
def display_frame(frame): def display_frame(frame):
......
...@@ -11,8 +11,12 @@ def load_database(file_location: str): ...@@ -11,8 +11,12 @@ def load_database(file_location: str):
:param file_location: String location of file to be loaded. :param file_location: String location of file to be loaded.
:return: The loaded pickle database. :return: The loaded pickle database.
""" """
with open(file_location, "rb") as database: file_content = {}
file_content = pickle.load(database) try:
with open(file_location, "rb") as database:
file_content = pickle.load(database)
except (FileNotFoundError, EOFError):
file_content = {}
return file_content return file_content
...@@ -23,5 +27,6 @@ def write_database(output_file: str, database_content: dict) -> None: ...@@ -23,5 +27,6 @@ def write_database(output_file: str, database_content: dict) -> None:
:param database_content: The database content to be written to the file. :param database_content: The database content to be written to the file.
:return: None :return: None
""" """
with open(output_file, "wb") as output: if output_file and database_content and database_content is not None:
output.write(pickle.dump(database_content, output)) with open(output_file, "wb") as output:
output.write(pickle.dump(database_content, output))
""" """
Creates a facial recognition profile for a new user. Creates a facial recognition profile for a new user.
""" """
import common
import data_handler
from common import USER_IDS_KEY, start_video_stream from common import USER_IDS_KEY, start_video_stream
import face_recognition import face_recognition
import cv2 import cv2
import os
def process_image(image, encoding_model: str = "hog"): def process_image(image, encoding_model: str = "hog"):
...@@ -21,14 +24,48 @@ def process_image(image, encoding_model: str = "hog"): ...@@ -21,14 +24,48 @@ def process_image(image, encoding_model: str = "hog"):
return face_recognition.face_encodings(image_rgb, boxes) return face_recognition.face_encodings(image_rgb, boxes)
def register_user(user_id: str, encoding_model="hog", image_flip: int = None, video_source=None): def register_user(user_id: str, dataset_dir: str, encoding_model="hog", database_loc: str = common.DATABASE_LOC,
show_output: bool = False):
""" """
Function for registering a new user using the given video source. If video source isn't provided, then the camera 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. on id 0 is used.
:param user_id: The user id for the user that is being registered. :param user_id: The user id for the user that is being registered.
:param dataset_dir: The directory location of pictures for the user.
:param encoding_model: The type of encoding model. Must be either "hog" or "cnn". HOG is faster, CNN is more thorough. :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 :param database_loc: Location of the pickle file database.
:return: Encoded face that was detected, or None if no face was detected or if there was another error. :return: Encoded face that was detected, or None if no face was detected or if there was another error.
""" """
if video_source is None: processed_images = []
video_source = start_video_stream(0) for (i, filename) in enumerate(os.listdir(dataset_dir)):
full_path = os.path.join(dataset_dir, filename)
# Might want to check file validity here at some point, but won't for now.
image = cv2.imread(full_path)
if image is not None:
processed = process_image(image, encoding_model=encoding_model)
if show_output:
print(f"Processing image {i + 1}")
if processed:
processed_images.append(processed)
if len(processed_images) > 0: # Only do things if we actually have stuff to add.
user_info = data_handler.load_database(database_loc)
if user_id not in user_info.keys():
user_info[user_id] = []
user_info[user_id].extend(processed_images)
data_handler.write_database(database_loc, user_info)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Facial Registration Options")
parser.add_argument("user_id", type=str, help="User ID to be used")
parser.add_argument("--encodings", "-e", type=str, help="File location to output encodings.", required=True,
default="./encodings.pickle")
parser.add_argument("--dataset", "-d", type=str, help="Directory location of the dataset images.", required=True)
parser.add_argument("--model", "-m", type=str, help="Type of encoding method, either \"hog\" or \"cnn\". HOG is "
"faster, CNN is more accurate.", required=False,
default="hog", choices=["cnn", "hog"])
args = parser.parse_args()
register_user(user_id=args.user_id, dataset_dir=args.dataset, encoding_model=args.model,
database_loc=args.encodings, show_output=True)
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment