wyzely-detect/wyzely_detect/__main__.py

# import face_recognition
from pathlib import Path

import cv2

# import hjson as json
import torch
from ultralytics import YOLO

from .utils import notify, utils
from .utils.cli_args import argparser

DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
args = None

objects_and_peoples = {
    "objects": {},
    "peoples": {},
}


def main():
    global objects_and_peoples
    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
    # https://github.com/ultralytics/ultralytics/issues/3084#issuecomment-1732433168
    # Currently, I have been unable to set up Poetry to use GPU for Torch
    for i in range(torch.cuda.device_count()):
        print(torch.cuda.get_device_properties(i).name)
    if torch.cuda.is_available():
        torch.cuda.set_device(0)
        print("Set CUDA device")
    else:
        print("No CUDA device available, using CPU")

    model = YOLO("yolov8n.pt")

    # Depending on if the user wants to use a stream or a capture device,
    # Set the video capture to the appropriate source
    if args.url:
        video_capture = cv2.VideoCapture(args.url)
    else:
        video_capture = cv2.VideoCapture(args.capture_device)

    # 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
    video_capture.set(cv2.CAP_PROP_BUFFERSIZE, 1)

    # Print the resolution of the video
    print(
        f"Video resolution: {video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)}x{video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)}"  # noqa: E501
    )

    print("Beginning video capture...")
    while True:
        # Grab a single frame of video
        ret, frame = video_capture.read()
        # Only process every other frame of video to save time
        # Resize frame of video to a smaller size for faster recognition processing
        run_frame = cv2.resize(frame, (0, 0), fx=args.run_scale, fy=args.run_scale)
        # view_frame = cv2.resize(frame, (0, 0), fx=args.view_scale, fy=args.view_scale)

        results = model(run_frame, verbose=False)

        path_to_faces = Path(args.faces_directory)
        path_to_faces_exists = path_to_faces.is_dir()

        for i, r in enumerate(results):
            # list of dicts with each dict containing a label, x1, y1, x2, y2
            plot_boxes = []

            # The following is stuff for people
            # This is still in the for loop as each result, no matter if anything is detected, will be present.
            # Thus, there will always be one result (r)

            # Only run if path_to_faces exists
            # 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
                ):
                    plot_boxes.append(face_details)
                    objects_and_peoples = notify.thing_detected(
                        thing_name=face_details["label"],
                        objects_and_peoples=objects_and_peoples,
                        detection_type="peoples",
                        detection_window=args.detection_window,
                        detection_duration=args.detection_duration,
                        notification_window=args.notification_window,
                        ntfy_url=args.ntfy_url,
                    )

            # The following is stuff for objects
            # Setup dictionary of object names
            if (
                objects_and_peoples["objects"] == {}
                or objects_and_peoples["objects"] is None
            ):
                for name in r.names.values():
                    objects_and_peoples["objects"][name] = {
                        "last_detection_time": None,
                        "detection_duration": None,
                        # "first_detection_time": None,
                        "last_notification_time": None,
                    }
                # Also, make sure that the objects to detect are in the list of objects_and_peoples
                # If it isn't, print a warning
                for obj in args.detect_object:
                    if obj not in objects_and_peoples:
                        print(
                            f"Warning: {obj} is not in the list of objects the model can detect!"
                        )

            for box in r.boxes:
                # Get the name of the object
                class_id = r.names[box.cls[0].item()]
                # Get the coordinates of the object
                cords = box.xyxy[0].tolist()
                cords = [round(x) for x in cords]
                # Get the confidence
                conf = round(box.conf[0].item(), 2)
                # Print it out, adding a spacer between each object
                # print("Object type:", class_id)
                # print("Coordinates:", cords)
                # print("Probability:", conf)
                # print("---")

                # Now do stuff (if conf > 0.5)
                if conf < args.object_confidence_threshold or (
                    class_id not in args.detect_object and args.detect_object != []
                ):
                    # If the confidence is too low
                    # or if the object is not in the list of objects to detect and the list of objects to detect is not empty
                    # then skip this iteration
                    continue

                # Add the object to the list of objects to plot
                plot_boxes.append(
                    {
                        "label": class_id,
                        "x1": cords[0],
                        "y1": cords[1],
                        "x2": cords[2],
                        "y2": cords[3],
                    }
                )

                objects_and_peoples = notify.thing_detected(
                    thing_name=class_id,
                    objects_and_peoples=objects_and_peoples,
                    detection_type="objects",
                    detection_window=args.detection_window,
                    detection_duration=args.detection_duration,
                    notification_window=args.notification_window,
                    ntfy_url=args.ntfy_url,
                )

            # To debug plotting, use r.plot() to cross reference the bounding boxes drawn by the plot_label() and r.plot()
            frame_to_show = utils.plot_label(
                boxes=plot_boxes,
                full_frame=frame,
                # full_frame=r.plot(),
                run_scale=args.run_scale,
                view_scale=args.view_scale,
            )

            # Display the resulting frame
            # cv2.imshow("", r)
            if not args.no_display:
                cv2.imshow(f"Video{i}", frame_to_show)

        # Hit 'q' on the keyboard to quit!
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    # Release handle to the webcam
    print("Releasing video capture")
    video_capture.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()