(untested) - Added args for multiple sources
Could not test due to dependency problems on Windows
This commit is contained in:
parent
5cc5e04642
commit
d3c157df4d
|
@ -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
|
||||
|
|
|
@ -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",
|
||||
|
|
Loading…
Reference in New Issue