(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
|
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
|
||||||
|
|
|
@ -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",
|
||||||
|
|
Loading…
Reference in New Issue