Add `--no-remove-representations`

This commit is contained in:
slashtechno 2023-10-27 11:33:05 -05:00
parent 32d523b727
commit 792a095782
Signed by: slashtechno
GPG Key ID: 8EC1D9D9286C2B17
3 changed files with 22 additions and 6 deletions

View File

@ -24,7 +24,6 @@ def main():
global args global args
# RUN_BY_COMPOSE = os.getenv("RUN_BY_COMPOSE") # Replace this with code to check for gpu # RUN_BY_COMPOSE = os.getenv("RUN_BY_COMPOSE") # Replace this with code to check for gpu
args = argparser.parse_args() args = argparser.parse_args()
# Check if a CUDA GPU is available. If it is, set it via torch. If not, set it to cpu # 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 # May be better to check every iteration, but this also works
if path_to_faces_exists: if path_to_faces_exists:
if face_details := utils.recognize_face( 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) plot_boxes.append(face_details)
objects_and_peoples = notify.thing_detected( objects_and_peoples = notify.thing_detected(

View File

@ -5,6 +5,7 @@ from pathlib import Path
argparser = None argparser = None
def set_argparse(): def set_argparse():
global argparser global argparser
@ -20,7 +21,6 @@ def set_argparse():
epilog=":)", epilog=":)",
) )
# One important thing to consider is that most function parameters are optional and have a default value # 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 # However, with argparse, those are never used since a argparse always passes something, even if it's None
@ -80,6 +80,15 @@ def set_argparse():
type=str, 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 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_<model>.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( argparser.add_argument(
"--detect-object", "--detect-object",
nargs="*", nargs="*",
@ -146,5 +155,6 @@ def set_argparse():
) )
# return argparser # return argparser
# This will run when this file is imported # This will run when this file is imported
set_argparse() set_argparse()

View File

@ -69,6 +69,7 @@ def recognize_face(
# opencv image # opencv image
run_frame: np.ndarray = None, run_frame: np.ndarray = None,
min_confidence: float = 0.3, min_confidence: float = 0.3,
no_remove_representations: bool = False,
) -> np.ndarray: ) -> np.ndarray:
""" """
Accepts a path to a directory of images of faces to be used as a refference 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 global first_face_try
# If it's the first time the function is being run, remove representations_arcface.pkl, if it exists # 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: try:
path_to_directory.joinpath("representations_arcface.pkl").unlink() path_to_directory.joinpath("representations_arcface.pkl").unlink()
print("Removing representations_arcface.pkl") print("Removing representations_arcface.pkl")
except FileNotFoundError: except FileNotFoundError:
print("representations_arcface.pkl does not exist") print("representations_arcface.pkl does not exist")
first_face_try = False 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 # 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 # It seems face_dataframes is empty if the face database (directory) doesn't exist. Seems to work if it's empty though