From d3c157df4d2337512dfc33992187f22ace710fea Mon Sep 17 00:00:00 2001
From: slashtechno <77907286+slashtechno@users.noreply.github.com>
Date: Sat, 10 Feb 2024 21:16:11 -0600
Subject: [PATCH] (untested) - Added args for multiple sources

Could not test due to dependency problems on Windows
---
 wyzely_detect/__main__.py       | 18 +++++++++++++-----
 wyzely_detect/utils/cli_args.py | 15 +++++++++------
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/wyzely_detect/__main__.py b/wyzely_detect/__main__.py
index 5db34e5..19ecf3b 100644
--- a/wyzely_detect/__main__.py
+++ b/wyzely_detect/__main__.py
@@ -2,6 +2,8 @@
 from pathlib import Path
 import cv2
 
+from PrettyTable import PrettyTable
+
 # import hjson as json
 import torch
 from ultralytics import YOLO
@@ -45,21 +47,27 @@ def main():
 
     # Depending on if the user wants to use a stream or a capture device,
     # Set the video capture to the appropriate source
-    if args.rtsp_url is not None:
-        video_capture = cv2.VideoCapture(args.rtsp_url)
+    if not args.rtsp_url and not args.capture_device:
+        print("No stream or capture device set, defaulting to capture device 0")
+        args.capture_device = [0]
     else:
-        video_capture = cv2.VideoCapture(args.capture_device)
+        video_sources = {
+            "streams": [cv2.VideoCapture(url) for url in args.rtsp_url],
+            "devices": [cv2.VideoCapture(device) for device in 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)
+    # Only do this for streams
+    for stream in video_sources["streams"]:
+        stream.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
     print("Beginning video capture...")
     while True:
         # Grab a single frame of video
diff --git a/wyzely_detect/utils/cli_args.py b/wyzely_detect/utils/cli_args.py
index 0f9e1ef..6337c76 100644
--- a/wyzely_detect/utils/cli_args.py
+++ b/wyzely_detect/utils/cli_args.py
@@ -29,17 +29,19 @@ def set_argparse():
     stream_source = video_options.add_mutually_exclusive_group()
     stream_source.add_argument(
         "--rtsp-url",
-        default=os.environ["RTSP_URL"]
-        if "RTSP_URL" in os.environ and os.environ["RTSP_URL"] != ""
-        else None,  # noqa: E501
+        action="append",
+        # 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
+        default=[os.environ["RTSP_URL"]] if "RTSP_URL" in os.environ and os.environ["RTSP_URL"] != "" else [],
         type=str,
         help="RTSP camera URL",
     )
     stream_source.add_argument(
         "--capture-device",
-        default=os.environ["CAPTURE_DEVICE"]
-        if "CAPTURE_DEVICE" in os.environ and os.environ["CAPTURE_DEVICE"] != ""
-        else 0,  # noqa: E501
+        action="append",
+        # 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
+        default=[int(os.environ["CAPTURE_DEVICE"])] if "CAPTURE_DEVICE" in os.environ and os.environ["CAPTURE_DEVICE"] != "" else [],
         type=int,
         help="Capture device number",
     )
@@ -155,6 +157,7 @@ def set_argparse():
     object_detection.add_argument(
         "--detect-object",
         action="append",
+        # Stuff is appended to default, as far as I can tell
         default=[],
         type=str,
         help="The object(s) to detect. Must be something the model is trained to detect",