(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 from pathlib import Path
import cv2 import cv2
from PrettyTable import PrettyTable
# import hjson as json # import hjson as json
import torch import torch
from ultralytics import YOLO from ultralytics import YOLO
@ -45,21 +47,27 @@ def main():
# Depending on if the user wants to use a stream or a capture device, # Depending on if the user wants to use a stream or a capture device,
# Set the video capture to the appropriate source # Set the video capture to the appropriate source
if args.rtsp_url is not None: if not args.rtsp_url and not args.capture_device:
video_capture = cv2.VideoCapture(args.rtsp_url) print("No stream or capture device set, defaulting to capture device 0")
args.capture_device = [0]
else: 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 # Eliminate lag by setting the buffer size to 1
# This makes it so that the video capture will only grab the most recent frame # This makes it so that the video capture will only grab the most recent frame
# However, this means that the video may be choppy # 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 the resolution of the video
print( print(
f"Video resolution: {video_capture.get(cv2.CAP_PROP_FRAME_WIDTH)}x{video_capture.get(cv2.CAP_PROP_FRAME_HEIGHT)}" # noqa: E501 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...") print("Beginning video capture...")
while True: while True:
# Grab a single frame of video # 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 = video_options.add_mutually_exclusive_group()
stream_source.add_argument( stream_source.add_argument(
"--rtsp-url", "--rtsp-url",
default=os.environ["RTSP_URL"] action="append",
if "RTSP_URL" in os.environ and os.environ["RTSP_URL"] != "" # If RTSP_URL is in the environment, use it, otherwise just use a blank list
else None, # noqa: E501 # 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, type=str,
help="RTSP camera URL", help="RTSP camera URL",
) )
stream_source.add_argument( stream_source.add_argument(
"--capture-device", "--capture-device",
default=os.environ["CAPTURE_DEVICE"] action="append",
if "CAPTURE_DEVICE" in os.environ and os.environ["CAPTURE_DEVICE"] != "" # If CAPTURE_DEVICE is in the environment, use it, otherwise just use a blank list
else 0, # noqa: E501 # 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, type=int,
help="Capture device number", help="Capture device number",
) )
@ -155,6 +157,7 @@ def set_argparse():
object_detection.add_argument( object_detection.add_argument(
"--detect-object", "--detect-object",
action="append", action="append",
# Stuff is appended to default, as far as I can tell
default=[], default=[],
type=str, type=str,
help="The object(s) to detect. Must be something the model is trained to detect", help="The object(s) to detect. Must be something the model is trained to detect",