Commit abfdef90 authored by Kyle Anderson's avatar Kyle Anderson

More small improvements

Just making refactorings and adding methods that I know I will need later.
parent 641f70d2
...@@ -26,8 +26,8 @@ def register_user(base_dir: str = os.getcwd()) -> None: ...@@ -26,8 +26,8 @@ def register_user(base_dir: str = os.getcwd()) -> None:
:param base_dir: The base directory, from which all default paths are relative. Defaults to current working directory. :param base_dir: The base directory, from which all default paths are relative. Defaults to current working directory.
:return: None :return: None
""" """
register.register_users_and_save(directory_location=os.path.join(base_dir, common.DATABASE_LOC), register.register_users_from_dir_and_save(directory_location=os.path.join(base_dir, common.DATABASE_LOC),
delete_images_on_complete=True, overwrite_data=False) delete_images_on_complete=True, overwrite_data=False)
if __name__ == "__main__": if __name__ == "__main__":
......
...@@ -36,22 +36,20 @@ def delete_file(file_path: str) -> None: ...@@ -36,22 +36,20 @@ def delete_file(file_path: str) -> None:
os.unlink(file_path) os.unlink(file_path)
def register_user(user_id: str, dataset_dir: str, encoding_model="hog", def register_user(user_id: str, images: list, encoding_model="hog",
show_output: bool = False, delete_on_processed: bool = False): 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 images of that user.
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 images: Images of the user. Should only contain the user, and no other faces.
: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 show_output: True to print console output for progress, false otherwise. :param show_output: True to print console output for progress, false otherwise.
:param delete_on_processed: True to delete the image file after processing it, false otherwise. :param delete_on_processed: True to delete the image file after processing it, false otherwise.
: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.
""" """
processed_images = [] processed_images = []
for (i, filename) in enumerate(impaths.list_images(dataset_dir)): for (i, image) in enumerate(images):
# Might want to check file validity here at some point, but won't for now. # Might want to check file validity here at some point, but won't for now.
image = cv2.imread(filename)
if image is not None: if image is not None:
if show_output: if show_output:
print(f"Processing image {i + 1} for user {user_id}") print(f"Processing image {i + 1} for user {user_id}")
...@@ -59,13 +57,52 @@ def register_user(user_id: str, dataset_dir: str, encoding_model="hog", ...@@ -59,13 +57,52 @@ def register_user(user_id: str, dataset_dir: str, encoding_model="hog",
if processed: if processed:
processed_images.extend(processed) processed_images.extend(processed)
# Delete after we're done if we're supposed to.
if delete_on_processed:
delete_file(filename)
return {user_id: processed_images} if len(processed_images) > 0 else None return {user_id: processed_images} if len(processed_images) > 0 else None
def register_user_from_directory(user_id: str, dataset_dir: str, encoding_model="hog",
show_output: bool = False, delete_on_processed: bool = False):
"""
Function for registering a new user using the directory location of images for that user.
: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 show_output: True to print console output for progress, false otherwise.
:param delete_on_processed: True to delete the image file after processing it, false otherwise.
:return: Encoded face that was detected, or None if no face was detected or if there was another error.
"""
images = []
filenames = []
for filename in impaths.list_images(dataset_dir):
filenames.append(filename)
images.append(cv2.imread(filename))
processed = register_user(user_id, images, encoding_model=encoding_model, show_output=show_output)
if delete_on_processed:
for filename in filenames:
delete_file(filename)
return processed
def register_users(user_dataset: dict, encoding_model: str = "hog", show_output: bool = False):
"""
Registers all the given users from their dataset images.
:param user_dataset: Dictionary of { "user_id" : [image1, image2, ...] , ... } of images of each user.
:param encoding_model: The type of encoding model to use.
:param show_output: True to print progress output, false otherwise.
:return: The dictionary of registered users in the given directory.
"""
total_dict = {}
for user_id, dataset in user_dataset.items():
user_dict = register_user(user_id, dataset, encoding_model=encoding_model, show_output=show_output)
if user_dict is not None:
total_dict.update(user_dict)
return total_dict if len(total_dict) > 0 else None
def register_users_in_dir(directory_location: str, encoding_model: str = "hog", delete_images_on_complete: bool = False, def register_users_in_dir(directory_location: str, encoding_model: str = "hog", delete_images_on_complete: bool = False,
show_output: bool = False): show_output: bool = False):
""" """
...@@ -80,21 +117,18 @@ def register_users_in_dir(directory_location: str, encoding_model: str = "hog", ...@@ -80,21 +117,18 @@ def register_users_in_dir(directory_location: str, encoding_model: str = "hog",
for directory in next(os.walk(directory_location))[1]: for directory in next(os.walk(directory_location))[1]:
total_directory = os.path.join(directory_location, directory) total_directory = os.path.join(directory_location, directory)
# Using the directory name as the user_id as well. # Using the directory name as the user_id as well.
user_dict = register_user(directory, total_directory, encoding_model=encoding_model, show_output=show_output, user_dict = register_user_from_directory(directory, total_directory, encoding_model=encoding_model,
delete_on_processed=delete_images_on_complete) show_output=show_output,
delete_on_processed=delete_images_on_complete)
if user_dict is not None: if user_dict is not None:
total_dict.update(user_dict) total_dict.update(user_dict)
return total_dict if len(total_dict) > 0 else None return total_dict if len(total_dict) > 0 else None
def register_users_and_save(directory_location: str = common.DATASET_DIR, def save_processed_users(processed_users,
database_location: str = common.DATABASE_LOC, encoding_model="hog", database_location: str = common.DATABASE_LOC,
delete_images_on_complete: bool = True, show_output: bool = False, overwrite_data: bool = False) -> None:
overwrite_data: bool = False):
processed_users = register_users_in_dir(directory_location, encoding_model=encoding_model,
delete_images_on_complete=delete_images_on_complete,
show_output=show_output)
database = data_handler.load_database(database_location) if not overwrite_data else {} database = data_handler.load_database(database_location) if not overwrite_data else {}
if processed_users is not None: if processed_users is not None:
for user_id, encodings in processed_users.items(): for user_id, encodings in processed_users.items():
...@@ -104,6 +138,26 @@ def register_users_and_save(directory_location: str = common.DATASET_DIR, ...@@ -104,6 +138,26 @@ def register_users_and_save(directory_location: str = common.DATASET_DIR,
data_handler.write_database(database_location, database) data_handler.write_database(database_location, database)
def register_users_from_dir_and_save(directory_location: str = common.DATASET_DIR,
database_location: str = common.DATABASE_LOC, encoding_model="hog",
delete_images_on_complete: bool = True, show_output: bool = False,
overwrite_data: bool = False) -> None:
"""
Registers the users in the given directory location and saves the user encodings to the database location.
:param directory_location: The location of the directory to crawl for dataset images.
:param database_location: The location of the database pickle file.
:param encoding_model: The encoding model which should be used. Hog is faster, cnn more reliable.
:param delete_images_on_complete: True to delete the images that are processed on completion, false otherwise.
:param show_output: True to show progress output to the console, false otherwise.
:param overwrite_data: True to overwrite the data in the database, false to append to it.
:return: None
"""
processed_users = register_users_in_dir(directory_location, encoding_model=encoding_model,
delete_images_on_complete=delete_images_on_complete,
show_output=show_output)
save_processed_users(processed_users, database_location=database_location, overwrite_data=overwrite_data)
if __name__ == "__main__": if __name__ == "__main__":
import argparse import argparse
...@@ -125,5 +179,5 @@ if __name__ == "__main__": ...@@ -125,5 +179,5 @@ if __name__ == "__main__":
if args.model is not None: if args.model is not None:
args_dict["encoding_model"] = args.model args_dict["encoding_model"] = args.model
register_users_and_save(**args_dict, show_output=True, delete_images_on_complete=False, register_users_from_dir_and_save(**args_dict, show_output=True, delete_images_on_complete=False,
overwrite_data=args.overwrite) overwrite_data=args.overwrite)
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