Fix various minor issues and format code

This commit is contained in:
slashtechno 2024-02-16 13:01:02 -06:00
parent b48edef250
commit c7d488d993
Signed by: slashtechno
GPG Key ID: 8EC1D9D9286C2B17
3 changed files with 86 additions and 92 deletions

View File

@ -1,7 +1,7 @@
# import face_recognition # import face_recognition
from pathlib import Path from pathlib import Path
import cv2 import cv2
import sys
from prettytable import PrettyTable from prettytable import PrettyTable
# import hjson as json # import hjson as json
@ -19,14 +19,13 @@ def main():
global objects_and_peoples global objects_and_peoples
global args global args
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
# https://github.com/ultralytics/ultralytics/issues/3084#issuecomment-1732433168 # https://github.com/ultralytics/ultralytics/issues/3084#issuecomment-1732433168
# Currently, I have been unable to set up Poetry to use GPU for Torch # Currently, I have been unable to set up Poetry to use GPU for Torch
for i in range(torch.cuda.device_count()): for i in range(torch.cuda.device_count()):
print(f'Using {torch.cuda.get_device_properties(i).name} for pytorch') print(f"Using {torch.cuda.get_device_properties(i).name} for pytorch")
if torch.cuda.is_available(): if torch.cuda.is_available():
torch.cuda.set_device(0) torch.cuda.set_device(0)
print("Set CUDA device") print("Set CUDA device")
@ -37,9 +36,10 @@ def main():
if args.force_disable_tensorflow_gpu: if args.force_disable_tensorflow_gpu:
print("Forcing tensorflow to use CPU") print("Forcing tensorflow to use CPU")
import tensorflow as tf import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')
if tf.config.experimental.list_logical_devices('GPU'): tf.config.set_visible_devices([], "GPU")
print('GPU disabled unsuccessfully') if tf.config.experimental.list_logical_devices("GPU"):
print("GPU disabled unsuccessfully")
else: else:
print("GPU disabled successfully") print("GPU disabled successfully")
@ -49,9 +49,7 @@ def main():
# Set the video capture to the appropriate source # Set the video capture to the appropriate source
if not args.rtsp_url and not args.capture_device: if not args.rtsp_url and not args.capture_device:
print("No stream or capture device set, defaulting to capture device 0") print("No stream or capture device set, defaulting to capture device 0")
video_sources = { video_sources = {"devices": [cv2.VideoCapture(0)]}
"devices": [cv2.VideoCapture(0)]
}
else: else:
video_sources = { video_sources = {
"streams": [cv2.VideoCapture(url) for url in args.rtsp_url], "streams": [cv2.VideoCapture(url) for url in args.rtsp_url],
@ -84,17 +82,22 @@ def main():
pretty_table = PrettyTable(field_names=["Source Type", "Resolution"]) pretty_table = PrettyTable(field_names=["Source Type", "Resolution"])
for source_type, sources in video_sources.items(): for source_type, sources in video_sources.items():
for source in sources: for source in sources:
if source.get(cv2.CAP_PROP_FRAME_WIDTH) == 0 or source.get(cv2.CAP_PROP_FRAME_HEIGHT) == 0: if (
source.get(cv2.CAP_PROP_FRAME_WIDTH) == 0
or source.get(cv2.CAP_PROP_FRAME_HEIGHT) == 0
):
message = "Capture for a source failed as resolution is 0x0.\n" message = "Capture for a source failed as resolution is 0x0.\n"
if source_type == "streams": if source_type == "streams":
message += "Check if the stream URL is correct and if the stream is online." message += "Check if the stream URL is correct and if the stream is online."
else: else:
message += "Check if the capture device is connected, working, and not in use by another program." message += "Check if the capture device is connected, working, and not in use by another program."
print(message) print(message)
# Maybe use os.exit() instead? sys.exit(1)
exit(1)
pretty_table.add_row( pretty_table.add_row(
[source_type, f"{source.get(cv2.CAP_PROP_FRAME_WIDTH)}x{source.get(cv2.CAP_PROP_FRAME_HEIGHT)}"] [
source_type,
f"{source.get(cv2.CAP_PROP_FRAME_WIDTH)}x{source.get(cv2.CAP_PROP_FRAME_HEIGHT)}",
]
) )
print(pretty_table) print(pretty_table)
@ -108,25 +111,23 @@ def main():
frames.extend([source.read()[1] for source in list_of_sources]) frames.extend([source.read()[1] for source in list_of_sources])
frames_to_show = [] frames_to_show = []
for frame in frames: for frame in frames:
frames_to_show.append(utils.process_footage( frames_to_show.append(
utils.process_footage(
frame=frame, frame=frame,
run_scale=args.run_scale, run_scale=args.run_scale,
view_scale=args.view_scale, view_scale=args.view_scale,
faces_directory=Path(args.faces_directory), faces_directory=Path(args.faces_directory),
face_confidence_threshold=args.face_confidence_threshold, face_confidence_threshold=args.face_confidence_threshold,
no_remove_representations=args.no_remove_representations, no_remove_representations=args.no_remove_representations,
detection_window=args.detection_window, detection_window=args.detection_window,
detection_duration=args.detection_duration, detection_duration=args.detection_duration,
notification_window=args.notification_window, notification_window=args.notification_window,
ntfy_url=args.ntfy_url, ntfy_url=args.ntfy_url,
model=model, model=model,
detect_object=args.detect_object, detect_object=args.detect_object,
object_confidence_threshold=args.object_confidence_threshold, object_confidence_threshold=args.object_confidence_threshold,
)) )
)
# Display the resulting frame # Display the resulting frame
# TODO: When multi-camera support is added, this needs to be changed to allow all feeds # TODO: When multi-camera support is added, this needs to be changed to allow all feeds
if not args.no_display: if not args.no_display:

View File

@ -15,16 +15,14 @@ def set_argparse():
else: else:
print("No .env file found") print("No .env file found")
# 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
argparser = argparse.ArgumentParser( argparser = argparse.ArgumentParser(
prog="Wyzely Detect", prog="Wyzely Detect",
description="Recognize faces/objects in a video stream (from a webcam or a security camera) and send notifications to your devices", # noqa: E501 description="Recognize faces/objects in a video stream (from a webcam or a security camera) and send notifications to your devices", # noqa: E501
epilog="For env bool options, setting them to anything except for an empty string will enable them." epilog="For env bool options, setting them to anything except for an empty string will enable them.",
) )
video_options = argparser.add_argument_group("Video Options") video_options = argparser.add_argument_group("Video Options")
stream_source = video_options.add_mutually_exclusive_group() stream_source = video_options.add_mutually_exclusive_group()
stream_source.add_argument( stream_source.add_argument(
@ -32,7 +30,9 @@ def set_argparse():
action="append", action="append",
# If RTSP_URL is in the environment, use it, otherwise just use a blank list # If RTSP_URL is in the environment, use it, otherwise just use a blank list
# This may cause problems down the road, but if it does, env for this can be removed # This may cause problems down the road, but if it does, env for this can be removed
default=[os.environ["RTSP_URL"]] if "RTSP_URL" in os.environ and os.environ["RTSP_URL"] != "" else [], default=[os.environ["RTSP_URL"]]
if "RTSP_URL" in os.environ and os.environ["RTSP_URL"] != ""
else [],
type=str, type=str,
help="RTSP camera URL", help="RTSP camera URL",
) )
@ -41,7 +41,9 @@ def set_argparse():
action="append", action="append",
# If CAPTURE_DEVICE is in the environment, use it, otherwise just use a blank list # If CAPTURE_DEVICE is in the environment, use it, otherwise just use a blank list
# If __main__.py detects that no capture device or remote stream is set, it will default to 0 # If __main__.py detects that no capture device or remote stream is set, it will default to 0
default=[int(os.environ["CAPTURE_DEVICE"])] if "CAPTURE_DEVICE" in os.environ and os.environ["CAPTURE_DEVICE"] != "" else [], default=[int(os.environ["CAPTURE_DEVICE"])]
if "CAPTURE_DEVICE" in os.environ and os.environ["CAPTURE_DEVICE"] != ""
else [],
type=int, type=int,
help="Capture device number", help="Capture device number",
) )
@ -77,8 +79,8 @@ def set_argparse():
help="Don't display the video feed", help="Don't display the video feed",
) )
video_options.add_argument( video_options.add_argument(
'-c', "-c",
'--force-disable-tensorflow-gpu', "--force-disable-tensorflow-gpu",
default=os.environ["FORCE_DISABLE_TENSORFLOW_GPU"] default=os.environ["FORCE_DISABLE_TENSORFLOW_GPU"]
if "FORCE_DISABLE_TENSORFLOW_GPU" in os.environ if "FORCE_DISABLE_TENSORFLOW_GPU" in os.environ
and os.environ["FORCE_DISABLE_TENSORFLOW_GPU"] != "" and os.environ["FORCE_DISABLE_TENSORFLOW_GPU"] != ""
@ -126,7 +128,6 @@ def set_argparse():
help="The time (seconds) before another notification can be sent", help="The time (seconds) before another notification can be sent",
) )
face_recognition = argparser.add_argument_group("Face Recognition options") face_recognition = argparser.add_argument_group("Face Recognition options")
face_recognition.add_argument( face_recognition.add_argument(
"--faces-directory", "--faces-directory",
@ -156,8 +157,6 @@ def set_argparse():
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 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
) )
object_detection = argparser.add_argument_group("Object Detection options") object_detection = argparser.add_argument_group("Object Detection options")
object_detection.add_argument( object_detection.add_argument(
"--detect-object", "--detect-object",
@ -174,8 +173,7 @@ def set_argparse():
and os.environ["OBJECT_CONFIDENCE_THRESHOLD"] != "" and os.environ["OBJECT_CONFIDENCE_THRESHOLD"] != ""
# I think this should always be a str so using lower shouldn't be a problem. # I think this should always be a str so using lower shouldn't be a problem.
# Also, if the first check fails the rest shouldn't be run # Also, if the first check fails the rest shouldn't be run
and os.environ["OBJECT_CONFIDENCE_THRESHOLD"].lower() != "false" and os.environ["OBJECT_CONFIDENCE_THRESHOLD"].lower() != "false" else 0.6,
else 0.6,
type=float, type=float,
help="The confidence threshold to use", help="The confidence threshold to use",
) )

View File

@ -2,8 +2,9 @@ import cv2
import os import os
import numpy as np import numpy as np
from pathlib import Path from pathlib import Path
# https://stackoverflow.com/a/42121886/18270659 # https://stackoverflow.com/a/42121886/18270659
os.environ['TF_CPP_MIN_LOG_LEVEL']='3' os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
from deepface import DeepFace # noqa: E402 from deepface import DeepFace # noqa: E402
@ -21,28 +22,23 @@ objects_and_peoples = {
def process_footage( def process_footage(
# Frame # Frame
frame: np.ndarray = None, frame: np.ndarray = None,
# scale # scale
run_scale: float = None, run_scale: float = None,
view_scale: float = None, view_scale: float = None,
# Face stuff # Face stuff
faces_directory: str = None, faces_directory: str = None,
face_confidence_threshold: float = None, face_confidence_threshold: float = None,
no_remove_representations: bool = False, no_remove_representations: bool = False,
# Timer stuff # Timer stuff
detection_window: int = None, detection_window: int = None,
detection_duration: int = None, detection_duration: int = None,
notification_window: int = None, notification_window: int = None,
ntfy_url: str = None, ntfy_url: str = None,
# Object stuff # Object stuff
# YOLO object # YOLO object
model=None, model=None,
detect_object: list = None, detect_object: list = None,
object_confidence_threshold = None object_confidence_threshold=None,
) -> np.ndarray: ) -> np.ndarray:
""" """
Takes in a frame and processes it Takes in a frame and processes it
@ -50,7 +46,6 @@ def process_footage(
global objects_and_peoples global objects_and_peoples
# Resize frame of video to a smaller size for faster recognition processing # Resize frame of video to a smaller size for faster recognition processing
run_frame = cv2.resize(frame, (0, 0), fx=run_scale, fy=run_scale) run_frame = cv2.resize(frame, (0, 0), fx=run_scale, fy=run_scale)
# view_frame = cv2.resize(frame, (0, 0), fx=args.view_scale, fy=args.view_scale) # view_frame = cv2.resize(frame, (0, 0), fx=args.view_scale, fy=args.view_scale)
@ -60,7 +55,7 @@ def process_footage(
path_to_faces = Path(faces_directory) path_to_faces = Path(faces_directory)
path_to_faces_exists = path_to_faces.is_dir() path_to_faces_exists = path_to_faces.is_dir()
for i, r in enumerate(results): for r in results:
# list of dicts with each dict containing a label, x1, y1, x2, y2 # list of dicts with each dict containing a label, x1, y1, x2, y2
plot_boxes = [] plot_boxes = []
@ -105,7 +100,8 @@ def process_footage(
# Also, make sure that the objects to detect are in the list of objects_and_peoples # Also, make sure that the objects to detect are in the list of objects_and_peoples
# If it isn't, print a warning # If it isn't, print a warning
for obj in detect_object: for obj in detect_object:
if obj not in objects_and_peoples["objects"].keys(): # .keys() shouldn't be needed
if obj not in objects_and_peoples["objects"]:
print( print(
f"Warning: {obj} is not in the list of objects the model can detect!" f"Warning: {obj} is not in the list of objects the model can detect!"
) )
@ -154,7 +150,6 @@ def process_footage(
ntfy_url=ntfy_url, ntfy_url=ntfy_url,
) )
# To debug plotting, use r.plot() to cross reference the bounding boxes drawn by the plot_label() and r.plot() # To debug plotting, use r.plot() to cross reference the bounding boxes drawn by the plot_label() and r.plot()
frame_to_show = plot_label( frame_to_show = plot_label(
boxes=plot_boxes, boxes=plot_boxes,
@ -238,7 +233,8 @@ def recognize_face(
Returns a single dictonary as currently only 1 face can be detected in each frame Returns a single dictonary as currently only 1 face can be detected in each frame
Cosine threshold is 0.3, so if the confidence is less than that, it will return None Cosine threshold is 0.3, so if the confidence is less than that, it will return None
dict contains the following keys: label, x1, y1, x2, y2 dict conta # Maybe use os.exit() instead?
ins the following keys: label, x1, y1, x2, y2
The directory should be structured as follows: The directory should be structured as follows:
faces/ faces/
name/ name/
@ -285,8 +281,11 @@ def recognize_face(
model_name="ArcFace", model_name="ArcFace",
detector_backend="opencv", detector_backend="opencv",
) )
"""
except (ValueError) as e: Example dataframe, for reference
identity (path to image) | source_x | source_y | source_w | source_h | VGG-Face_cosine (pretty much the confidence \_('_')_/)
"""
except ValueError as e:
if ( if (
str(e) str(e)
== "Face could not be detected. Please confirm that the picture is a face photo or consider to set enforce_detection param to False." # noqa: E501 == "Face could not be detected. Please confirm that the picture is a face photo or consider to set enforce_detection param to False." # noqa: E501
@ -295,7 +294,8 @@ def recognize_face(
return None return None
elif ( elif (
# Check if the error message contains "Validate .jpg or .png files exist in this path." # Check if the error message contains "Validate .jpg or .png files exist in this path."
"Validate .jpg or .png files exist in this path." in str(e) "Validate .jpg or .png files exist in this path."
in str(e)
): ):
# If a verbose/silent flag is added, this should be changed to print only if verbose is true # If a verbose/silent flag is added, this should be changed to print only if verbose is true
# print("No faces found in database") # print("No faces found in database")
@ -338,8 +338,3 @@ def recognize_face(
) )
return to_return return to_return
return None return None
"""
Example dataframe, for reference
identity (path to image) | source_x | source_y | source_w | source_h | VGG-Face_cosine (pretty much the confidence \_('_')_/)
"""