2023-10-02 01:56:40 +01:00
|
|
|
# import face_recognition
|
|
|
|
from pathlib import Path
|
2023-10-27 16:54:36 +01:00
|
|
|
import cv2
|
2024-02-16 19:01:02 +00:00
|
|
|
import sys
|
2024-02-11 20:27:05 +00:00
|
|
|
from prettytable import PrettyTable
|
2024-02-11 03:16:11 +00:00
|
|
|
|
2023-10-02 01:56:40 +01:00
|
|
|
# import hjson as json
|
|
|
|
import torch
|
|
|
|
from ultralytics import YOLO
|
|
|
|
|
2023-12-22 21:22:01 +00:00
|
|
|
from .utils import utils
|
2023-10-27 16:54:36 +01:00
|
|
|
from .utils.cli_args import argparser
|
2023-10-02 01:56:40 +01:00
|
|
|
|
|
|
|
DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
|
|
|
|
args = None
|
|
|
|
|
2023-10-05 03:03:11 +01:00
|
|
|
|
2023-10-02 01:56:40 +01:00
|
|
|
def main():
|
2024-02-16 19:01:02 +00:00
|
|
|
global args
|
2023-10-02 01:56:40 +01:00
|
|
|
|
|
|
|
args = argparser.parse_args()
|
|
|
|
|
2023-10-22 22:53:29 +01:00
|
|
|
# Check if a CUDA GPU is available. If it is, set it via torch. If not, set it to cpu
|
2023-10-02 01:56:40 +01:00
|
|
|
# https://github.com/ultralytics/ultralytics/issues/3084#issuecomment-1732433168
|
2023-10-05 03:03:11 +01:00
|
|
|
# Currently, I have been unable to set up Poetry to use GPU for Torch
|
2023-10-04 03:03:39 +01:00
|
|
|
for i in range(torch.cuda.device_count()):
|
2024-02-16 19:01:02 +00:00
|
|
|
print(f"Using {torch.cuda.get_device_properties(i).name} for pytorch")
|
2023-10-04 03:03:39 +01:00
|
|
|
if torch.cuda.is_available():
|
2023-10-02 01:56:40 +01:00
|
|
|
torch.cuda.set_device(0)
|
|
|
|
print("Set CUDA device")
|
2023-10-05 03:03:11 +01:00
|
|
|
else:
|
2023-10-02 01:56:40 +01:00
|
|
|
print("No CUDA device available, using CPU")
|
2023-12-03 03:23:42 +00:00
|
|
|
# Seems automatically, deepface (tensorflow) tried to use my GPU on Pop!_OS (I did not set up cudnn or anything)
|
|
|
|
# Not sure the best way, in Poetry, to manage GPU libraries so for now, just use CPU
|
|
|
|
if args.force_disable_tensorflow_gpu:
|
|
|
|
print("Forcing tensorflow to use CPU")
|
|
|
|
import tensorflow as tf
|
2024-02-16 19:01:02 +00:00
|
|
|
|
|
|
|
tf.config.set_visible_devices([], "GPU")
|
|
|
|
if tf.config.experimental.list_logical_devices("GPU"):
|
|
|
|
print("GPU disabled unsuccessfully")
|
2023-12-03 03:23:42 +00:00
|
|
|
else:
|
|
|
|
print("GPU disabled successfully")
|
2023-10-05 03:03:11 +01:00
|
|
|
|
2023-10-02 01:56:40 +01:00
|
|
|
model = YOLO("yolov8n.pt")
|
|
|
|
|
2023-10-06 00:40:53 +01:00
|
|
|
# Depending on if the user wants to use a stream or a capture device,
|
|
|
|
# Set the video capture to the appropriate source
|
2024-02-11 03:16:11 +00:00
|
|
|
if not args.rtsp_url and not args.capture_device:
|
|
|
|
print("No stream or capture device set, defaulting to capture device 0")
|
2024-02-16 19:01:02 +00:00
|
|
|
video_sources = {"devices": [cv2.VideoCapture(0)]}
|
2023-10-06 00:40:53 +01:00
|
|
|
else:
|
2024-02-11 03:16:11 +00:00
|
|
|
video_sources = {
|
|
|
|
"streams": [cv2.VideoCapture(url) for url in args.rtsp_url],
|
|
|
|
"devices": [cv2.VideoCapture(device) for device in args.capture_device],
|
|
|
|
}
|
2023-10-06 00:40:53 +01:00
|
|
|
|
2024-02-16 18:34:39 +00:00
|
|
|
if args.fake_second_source:
|
|
|
|
try:
|
|
|
|
video_sources["devices"].append(video_sources["devices"][0])
|
|
|
|
except KeyError:
|
|
|
|
print("No capture device to use as second source. Trying stream.")
|
|
|
|
try:
|
|
|
|
video_sources["devices"].append(video_sources["devices"][0])
|
|
|
|
except KeyError:
|
|
|
|
print("No stream to use as a second source")
|
|
|
|
# When the code tries to resize the nonexistent capture device 1, the program will fail
|
|
|
|
|
2023-10-02 01:56:40 +01:00
|
|
|
# Eliminate lag by setting the buffer size to 1
|
|
|
|
# This makes it so that the video capture will only grab the most recent frame
|
|
|
|
# However, this means that the video may be choppy
|
2024-02-11 03:16:11 +00:00
|
|
|
# Only do this for streams
|
2024-02-11 21:41:46 +00:00
|
|
|
try:
|
|
|
|
for stream in video_sources["streams"]:
|
|
|
|
stream.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|
|
|
# If there are no streams, this will throw a KeyError
|
|
|
|
except KeyError:
|
|
|
|
pass
|
|
|
|
|
|
|
|
# Print out the resolution of the video sources. Ideally, change this so the device ID/url is also printed
|
|
|
|
pretty_table = PrettyTable(field_names=["Source Type", "Resolution"])
|
|
|
|
for source_type, sources in video_sources.items():
|
|
|
|
for source in sources:
|
2024-02-16 19:01:02 +00:00
|
|
|
if (
|
|
|
|
source.get(cv2.CAP_PROP_FRAME_WIDTH) == 0
|
|
|
|
or source.get(cv2.CAP_PROP_FRAME_HEIGHT) == 0
|
|
|
|
):
|
2024-02-16 18:45:31 +00:00
|
|
|
message = "Capture for a source failed as resolution is 0x0.\n"
|
|
|
|
if source_type == "streams":
|
|
|
|
message += "Check if the stream URL is correct and if the stream is online."
|
|
|
|
else:
|
|
|
|
message += "Check if the capture device is connected, working, and not in use by another program."
|
|
|
|
print(message)
|
2024-02-16 19:01:02 +00:00
|
|
|
sys.exit(1)
|
2024-02-11 21:41:46 +00:00
|
|
|
pretty_table.add_row(
|
2024-02-16 19:01:02 +00:00
|
|
|
[
|
|
|
|
source_type,
|
|
|
|
f"{source.get(cv2.CAP_PROP_FRAME_WIDTH)}x{source.get(cv2.CAP_PROP_FRAME_HEIGHT)}",
|
|
|
|
]
|
2024-02-11 21:41:46 +00:00
|
|
|
)
|
|
|
|
print(pretty_table)
|
2023-10-02 01:56:40 +01:00
|
|
|
print("Beginning video capture...")
|
|
|
|
while True:
|
|
|
|
# Grab a single frame of video
|
2024-02-16 18:34:39 +00:00
|
|
|
frames = []
|
|
|
|
# frames = [source.read() for sources in video_sources.values() for source in sources]
|
|
|
|
for list_of_sources in video_sources.values():
|
|
|
|
frames.extend([source.read()[1] for source in list_of_sources])
|
|
|
|
frames_to_show = []
|
|
|
|
for frame in frames:
|
2024-02-16 19:01:02 +00:00
|
|
|
frames_to_show.append(
|
|
|
|
utils.process_footage(
|
|
|
|
frame=frame,
|
|
|
|
run_scale=args.run_scale,
|
|
|
|
view_scale=args.view_scale,
|
|
|
|
faces_directory=Path(args.faces_directory),
|
|
|
|
face_confidence_threshold=args.face_confidence_threshold,
|
|
|
|
no_remove_representations=args.no_remove_representations,
|
|
|
|
detection_window=args.detection_window,
|
|
|
|
detection_duration=args.detection_duration,
|
|
|
|
notification_window=args.notification_window,
|
|
|
|
ntfy_url=args.ntfy_url,
|
|
|
|
model=model,
|
|
|
|
detect_object=args.detect_object,
|
|
|
|
object_confidence_threshold=args.object_confidence_threshold,
|
|
|
|
)
|
|
|
|
)
|
2023-12-22 21:22:01 +00:00
|
|
|
# Display the resulting frame
|
|
|
|
# TODO: When multi-camera support is added, this needs to be changed to allow all feeds
|
|
|
|
if not args.no_display:
|
2024-02-16 19:01:02 +00:00
|
|
|
for i, frame_to_show in enumerate(frames_to_show):
|
2024-02-16 18:34:39 +00:00
|
|
|
cv2.imshow(f"Video {i}", frame_to_show)
|
2023-10-02 01:56:40 +01:00
|
|
|
|
|
|
|
# Hit 'q' on the keyboard to quit!
|
|
|
|
if cv2.waitKey(1) & 0xFF == ord("q"):
|
|
|
|
break
|
|
|
|
|
|
|
|
# Release handle to the webcam
|
|
|
|
print("Releasing video capture")
|
2024-02-16 18:34:39 +00:00
|
|
|
[source.release() for sources in video_sources.values() for source in sources]
|
2023-10-02 01:56:40 +01:00
|
|
|
cv2.destroyAllWindows()
|
|
|
|
|
2023-10-15 01:25:27 +01:00
|
|
|
|
2023-10-14 21:46:42 +01:00
|
|
|
if __name__ == "__main__":
|
|
|
|
main()
|