register_user.py 5.97 KB
Newer Older
1 2 3
"""
Creates a facial recognition profile for a new user.
"""
4 5 6 7
import os

import cv2
import face_recognition
Kyle Anderson's avatar
Kyle Anderson committed
8
from imutils import paths as impaths
9

Kyle Anderson's avatar
Kyle Anderson committed
10 11
import common
import data_handler
12 13 14 15 16 17 18 19 20 21 22 23 24 25


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.
26 27
    # Only want the first recognized face
    return face_recognition.face_encodings(image_rgb, [boxes[0]]) if boxes and len(boxes) > 0 else []
28 29


Kyle Anderson's avatar
Kyle Anderson committed
30 31 32 33 34 35 36 37 38 39 40
def delete_file(file_path: str) -> None:
    """
    Deletes the file at the given location.
    :param file_path: The path to the file.
    :return: None
    """
    os.unlink(file_path)


def register_user(user_id: str, dataset_dir: str, encoding_model="hog",
                  show_output: bool = False, delete_on_processed: bool = False):
41 42 43 44
    """
    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.
Kyle Anderson's avatar
Kyle Anderson committed
45
    :param dataset_dir: The directory location of pictures for the user.
46
    :param encoding_model: The type of encoding model. Must be either "hog" or "cnn". HOG is faster, CNN is more thorough.
Kyle Anderson's avatar
Kyle Anderson committed
47 48
    :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.
49 50
    :return: Encoded face that was detected, or None if no face was detected or if there was another error.
    """
Kyle Anderson's avatar
Kyle Anderson committed
51
    processed_images = []
Kyle Anderson's avatar
Kyle Anderson committed
52
    for (i, filename) in enumerate(impaths.list_images(dataset_dir)):
Kyle Anderson's avatar
Kyle Anderson committed
53
        # Might want to check file validity here at some point, but won't for now.
Kyle Anderson's avatar
Kyle Anderson committed
54
        image = cv2.imread(filename)
Kyle Anderson's avatar
Kyle Anderson committed
55 56
        if image is not None:
            if show_output:
Kyle Anderson's avatar
Kyle Anderson committed
57
                print(f"Processing image {i + 1} for user {user_id}")
58
            processed = process_image(image, encoding_model=encoding_model)
Kyle Anderson's avatar
Kyle Anderson committed
59
            if processed:
60
                processed_images.extend(processed)
Kyle Anderson's avatar
Kyle Anderson committed
61

Kyle Anderson's avatar
Kyle Anderson committed
62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
                # 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


def register_users_in_dir(directory_location: str, encoding_model: str = "hog", delete_images_on_complete: bool = False,
                          show_output: bool = False):
    """
    Registers all the users in a directory.
    :param directory_location:
    :param encoding_model: The type of encoding model to use.
    :param delete_images_on_complete: True to delete the images after processing them, false otherwise.
    :param show_output: True to print progress output, false otherwise.
    :return: The dictionary of registered users in the given directory.
    """
    total_dict = {}
    for directory in next(os.walk(directory_location))[1]:
        total_directory = os.path.join(directory_location, directory)
        # 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,
                                  delete_on_processed=delete_images_on_complete)
        if user_dict is not None:
            total_dict.update(user_dict)

    return total_dict if len(total_dict) > 0 else None


def register_users_and_save(directory_location: str = common.DATASET_DIR,
Kyle Anderson's avatar
Kyle Anderson committed
92
                            database_location: str = common.DATABASE_LOC, encoding_model="hog",
Kyle Anderson's avatar
Kyle Anderson committed
93 94 95 96 97 98 99 100 101 102 103 104
                            delete_images_on_complete: bool = True, show_output: bool = False,
                            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 {}
    if processed_users is not None:
        for user_id, encodings in processed_users.items():
            if user_id not in database:
                database[user_id] = []
            database[user_id].extend(encodings)
    data_handler.write_database(database_location, database)
Kyle Anderson's avatar
Kyle Anderson committed
105 106 107 108 109 110


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Facial Registration Options")
Kyle Anderson's avatar
Kyle Anderson committed
111 112
    parser.add_argument("--encodings", "-e", type=str, help="File location to output encodings.", default=None)
    parser.add_argument("--dataset", "-d", type=str, help="Directory location of the dataset images.", default=None)
Kyle Anderson's avatar
Kyle Anderson committed
113
    parser.add_argument("--model", "-m", type=str, help="Type of encoding method, either \"hog\" or \"cnn\". HOG is "
Kyle Anderson's avatar
Kyle Anderson committed
114 115 116 117
                                                        "faster, CNN is more accurate.", required=False, default=None,
                        choices=["cnn", "hog"])
    parser.add_argument("--overwrite", default=False, action="store_true",
                        help="Include this flag to overwrite the database, replacing its content.", required=False)
Kyle Anderson's avatar
Kyle Anderson committed
118
    args = parser.parse_args()
Kyle Anderson's avatar
Kyle Anderson committed
119 120 121 122 123 124 125 126 127 128 129

    args_dict = {}
    if args.encodings is not None:
        args_dict["database_location"] = args.encodings
    if args.dataset is not None:
        args_dict["directory_location"] = args.dataset
    if args.model is not None:
        args_dict["encoding_model"] = args.model

    register_users_and_save(**args_dict, show_output=True, delete_images_on_complete=False,
                            overwrite_data=args.overwrite)