diff --git a/wyzely_detect/__main__.py b/wyzely_detect/__main__.py index 3c46edd..72e4ea0 100644 --- a/wyzely_detect/__main__.py +++ b/wyzely_detect/__main__.py @@ -24,7 +24,6 @@ def main(): global args # RUN_BY_COMPOSE = os.getenv("RUN_BY_COMPOSE") # Replace this with code to check for gpu - args = argparser.parse_args() # Check if a CUDA GPU is available. If it is, set it via torch. If not, set it to cpu @@ -83,7 +82,10 @@ def main(): # May be better to check every iteration, but this also works if path_to_faces_exists: if face_details := utils.recognize_face( - path_to_directory=path_to_faces, run_frame=run_frame, min_confidence=args.face_confidence_threshold + path_to_directory=path_to_faces, + run_frame=run_frame, + min_confidence=args.face_confidence_threshold, + no_remove_representations=args.no_remove_representations, ): plot_boxes.append(face_details) objects_and_peoples = notify.thing_detected( diff --git a/wyzely_detect/utils/cli_args.py b/wyzely_detect/utils/cli_args.py index 5c064c6..b343057 100644 --- a/wyzely_detect/utils/cli_args.py +++ b/wyzely_detect/utils/cli_args.py @@ -5,6 +5,7 @@ from pathlib import Path argparser = None + def set_argparse(): global argparser @@ -20,10 +21,9 @@ def set_argparse(): epilog=":)", ) - # One important thing to consider is that most function parameters are optional and have a default value # However, with argparse, those are never used since a argparse always passes something, even if it's None - + argparser.add_argument( "--run-scale", # Set it to the env RUN_SCALE if it isn't blank, otherwise set it to 0.25 @@ -80,6 +80,15 @@ def set_argparse(): type=str, help="The directory to store the faces. Can either contain images or subdirectories with images, the latter being the preferred method", # noqa: E501 ) + argparser.add_argument( + "--no-remove-representations", + default=os.environ["NO_REMOVE_REPRESENTATIONS"] + if "NO_REMOVE_REPRESENTATIONS" in os.environ + and os.environ["NO_REMOVE_REPRESENTATIONS"] != "" + else False, + action="store_true", + help="Don't remove representations_.pkl at the start of the program. Greatly improves startup time, but doesn't take into account changes to the faces directory since it was created", # noqa: E501 + ) argparser.add_argument( "--detect-object", nargs="*", @@ -146,5 +155,6 @@ def set_argparse(): ) # return argparser + # This will run when this file is imported -set_argparse() \ No newline at end of file +set_argparse() diff --git a/wyzely_detect/utils/utils.py b/wyzely_detect/utils/utils.py index 17592e3..d7ffddf 100644 --- a/wyzely_detect/utils/utils.py +++ b/wyzely_detect/utils/utils.py @@ -69,6 +69,7 @@ def recognize_face( # opencv image run_frame: np.ndarray = None, min_confidence: float = 0.3, + no_remove_representations: bool = False, ) -> np.ndarray: """ Accepts a path to a directory of images of faces to be used as a refference @@ -95,13 +96,16 @@ def recognize_face( global first_face_try # If it's the first time the function is being run, remove representations_arcface.pkl, if it exists - if first_face_try: + if first_face_try and not no_remove_representations: try: path_to_directory.joinpath("representations_arcface.pkl").unlink() print("Removing representations_arcface.pkl") except FileNotFoundError: print("representations_arcface.pkl does not exist") first_face_try = False + elif first_face_try and no_remove_representations: + print("Not attempting to remove representations_arcface.pkl") + first_face_try = False # face_dataframes is a vanilla list of dataframes # It seems face_dataframes is empty if the face database (directory) doesn't exist. Seems to work if it's empty though