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
recognized_users = check_recognized_users(recognized_users_count)
if len(recognized_users) > 0:
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:
recognized_user = max(recognized_users_count,
......@@ -142,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=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)
if user:
print(f"Recognized user {user}.")
......@@ -4,6 +4,7 @@ import cv2
from imutils.video import VideoStream
USER_IDS_KEY = "user_ids"
DATABASE_LOC = "./dataset/faces.pickle"
def display_frame(frame):
......
......@@ -11,8 +11,12 @@ def load_database(file_location: str):
: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)
file_content = {}
try:
with open(file_location, "rb") as database:
file_content = pickle.load(database)
except (FileNotFoundError, EOFError):
file_content = {}
return file_content
......@@ -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.
:return: None
"""
with open(output_file, "wb") as output:
output.write(pickle.dump(database_content, output))
if output_file and database_content and database_content is not None:
with open(output_file, "wb") as output:
output.write(pickle.dump(database_content, output))
"""
Creates a facial recognition profile for a new user.
"""
import common
import data_handler
from common import USER_IDS_KEY, start_video_stream
import face_recognition
import cv2
import os
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)
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
on id 0 is used.
: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 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.
"""
if video_source is None:
video_source = start_video_stream(0)
processed_images = []
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