(untested) - Added args for multiple sources

Could not test due to dependency problems on Windows
This commit is contained in:
slashtechno 2024-02-10 21:16:11 -06:00
parent 5cc5e04642
commit d3c157df4d
Signed by: slashtechno
GPG Key ID: 8EC1D9D9286C2B17
2 changed files with 22 additions and 11 deletions

View File

@ -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

View File

@ -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",