xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

examples/time_in_zone: improve CLI argument handling (#2057) 

* refactor(time_in_zone): improve CLI argument handling
* Refactor examples with `jsonargparse` for improved CLI, rename `confidence` and `iou` arguments for consistency, and update documentation and requirements.
* fix(pre_commit): 🎨 auto format pre-commit hook

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Jirka Borovec 2026-02-03 20:49:55 +09:00 committed by GitHub
parent 1688e71ed9
commit 11f483b75e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 216 additions and 444 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
examples/time_in_zone/README.md
View File

@ -129,9 +129,10 @@ python inference_file_example.py \
--zone_configuration_path "data/checkout/config.json" \
--source_video_path "data/checkout/video.mp4" \
--model_id "yolov8x-640" \
--classes 0 \
--classes "[0]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
--iou_threshold 0.7 \
--roboflow_api_key "ROBOFLOWS_API_KEY"
```
https://github.com/roboflow/supervision/assets/26109316/d051cc8a-dd15-41d4-aa36-d38b86334c39
@ -141,9 +142,10 @@ python inference_file_example.py \
--zone_configuration_path "data/traffic/config.json" \
--source_video_path "data/traffic/video.mp4" \
--model_id "yolov8x-640" \
--classes 2 5 6 7 \
--classes "[2, 5, 6, 7]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
--iou_threshold 0.7 \
--roboflow_api_key "ROBOFLOWS_API_KEY"
```
https://github.com/roboflow/supervision/assets/26109316/5ec896d7-4b39-4426-8979-11e71666878b
@ -164,7 +166,7 @@ python inference_stream_example.py \
--zone_configuration_path "data/checkout/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--model_id "yolov8x-640" \
--classes 0 \
--classes "[0]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
@ -174,7 +176,7 @@ python inference_stream_example.py \
--zone_configuration_path "data/traffic/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--model_id "yolov8x-640" \
--classes 2 5 6 7 \
--classes "[2, 5, 6, 7]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
@ -200,7 +202,7 @@ python ultralytics_file_example.py \
--source_video_path "data/checkout/video.mp4" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 0 \
--classes "[0]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
@ -211,7 +213,7 @@ python ultralytics_file_example.py \
--source_video_path "data/traffic/video.mp4" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 2 5 6 7 \
--classes "[2, 5, 6, 7]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
@ -234,7 +236,7 @@ python ultralytics_stream_example.py \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 0 \
--classes "[0]" \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```

80
examples/time_in_zone/inference_file_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from inference import get_model
@ -16,14 +14,27 @@ LABEL_ANNOTATOR = sv.LabelAnnotator(
def main(
source_video_path: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: list[int],
source_video_path: str,
model_id: str = "yolov8s-640",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
roboflow_api_key: str = "",
) -> None:
model = get_model(model_id=model_id)
"""
Calculating detections dwell time in zones, using video file.
Args:
zone_configuration_path: Path to the zone configuration JSON file
source_video_path: Path to the source video file
model_id: Roboflow model ID
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
roboflow_api_key: Roboflow API key for accessing private models
"""
model = get_model(model_id=model_id, api_key=roboflow_api_key)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
frames_generator = sv.get_video_frames_generator(source_video_path)
@ -39,7 +50,9 @@ def main(
timers = [FPSBasedTimer(video_info.fps) for _ in zones]
for frame in frames_generator:
results = model.infer(frame, confidence=confidence, iou_threshold=iou)[0]
results = model.infer(
frame, confidence=confidence_threshold, iou_threshold=iou_threshold
)[0]
detections = sv.Detections.from_inference(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = tracker.update_with_detections(detections)
@ -78,50 +91,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using video file."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--source_video_path",
type=str,
required=True,
help="Path to the source video file.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
source_video_path=args.source_video_path,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

80
examples/time_in_zone/inference_naive_stream_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from inference import get_model
@ -16,14 +14,27 @@ LABEL_ANNOTATOR = sv.LabelAnnotator(
def main(
rtsp_url: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: list[int],
rtsp_url: str,
model_id: str = "yolov8s-640",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
roboflow_api_key: str = "",
) -> None:
model = get_model(model_id=model_id)
"""
Calculating detections dwell time in zones, using RTSP stream.
Args:
zone_configuration_path: Path to the zone configuration JSON file
rtsp_url: Complete RTSP URL for the video stream
model_id: Roboflow model ID
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
roboflow_api_key: Roboflow API key for accessing private models
"""
model = get_model(model_id=model_id, api_key=roboflow_api_key)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
fps_monitor = sv.FPSMonitor()
@ -42,7 +53,9 @@ def main(
fps_monitor.tick()
fps = fps_monitor.fps
results = model.infer(frame, confidence=confidence, iou_threshold=iou)[0]
results = model.infer(
frame, confidence=confidence_threshold, iou_threshold=iou_threshold
)[0]
detections = sv.Detections.from_inference(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = tracker.update_with_detections(detections)
@ -88,50 +101,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

79
examples/time_in_zone/inference_stream_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from inference import InferencePipeline
@ -77,21 +75,35 @@ class CustomSink:
def main(
rtsp_url: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: list[int],
rtsp_url: str,
model_id: str = "yolov8s-640",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
roboflow_api_key: str = "",
) -> None:
"""
Calculating detections dwell time in zones, using RTSP stream.
Args:
zone_configuration_path: Path to the zone configuration JSON file
rtsp_url: Complete RTSP URL for the video stream
model_id: Roboflow model ID
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
roboflow_api_key: Roboflow API key for accessing private models
"""
sink = CustomSink(zone_configuration_path=zone_configuration_path, classes=classes)
pipeline = InferencePipeline.init(
model_id=model_id,
video_reference=rtsp_url,
on_prediction=sink.on_prediction,
confidence=confidence,
iou_threshold=iou,
confidence=confidence_threshold,
iou_threshold=iou_threshold,
api_key=roboflow_api_key,
)
pipeline.start()
@ -103,50 +115,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

3
examples/time_in_zone/requirements.txt
View File

@ -1,6 +1,7 @@
supervision
ultralytics
inference
inference[gaze,transformers,sam]
# https://github.com/pytube/pytube/issues/2044
# pytube
pytubefix
jsonargparse[signatures]

46
examples/time_in_zone/scripts/download_from_youtube.py
View File

@ -1,12 +1,25 @@
from __future__ import annotations
import argparse
import os
import ssl
from jsonargparse import auto_cli
from pytubefix import YouTube
def main(url: str, output_path: str | None, file_name: str | None) -> None:
def main(
url: str, output_path: str = "data/source", file_name: str = "video.mp4"
) -> None:
"""
Download a specific YouTube video by providing its URL.
Args:
url: The full URL of the YouTube video you wish to download.
output_path: Specifies the directory where the video will be saved.
file_name: Sets the name of the saved video file.
"""
ssl._create_default_https_context = ssl._create_unverified_context
yt = YouTube(url)
stream = yt.streams.get_highest_resolution()
@ -20,28 +33,7 @@ def main(url: str, output_path: str | None, file_name: str | None) -> None:
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Download a specific YouTube video by providing its URL."
)
parser.add_argument(
"--url",
type=str,
required=True,
help="The full URL of the YouTube video you wish to download.",
)
parser.add_argument(
"--output_path",
type=str,
default="data/source",
required=False,
help="Optional. Specifies the directory where the video will be saved.",
)
parser.add_argument(
"--file_name",
type=str,
default="video.mp4",
required=False,
help="Optional. Sets the name of the saved video file.",
)
args = parser.parse_args()
main(url=args.url, output_path=args.output_path, file_name=args.file_name)
from jsonargparse import auto_cli, set_parsing_settings
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

34
examples/time_in_zone/scripts/draw_zones.py
View File

@ -1,12 +1,12 @@
from __future__ import annotations
import argparse
import json
import os
from typing import Any
import cv2
import numpy as np
from jsonargparse import auto_cli
import supervision as sv
@ -126,6 +126,13 @@ def save_polygons_to_json(polygons, target_path):
def main(source_path: str, zone_configuration_path: str) -> None:
"""
Interactively draw polygons on images or video frames and save the annotations.
Args:
source_path: Path to the source image or video file for drawing polygons.
zone_configuration_path: Path where the polygon annotations saved as JSON file.
"""
global current_mouse_position
original_image = resolve_source(source_path=source_path)
if original_image is None:
@ -155,24 +162,7 @@ def main(source_path: str, zone_configuration_path: str) -> None:
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Interactively draw polygons on images or video frames and save "
"the annotations."
)
parser.add_argument(
"--source_path",
type=str,
required=True,
help="Path to the source image or video file for drawing polygons.",
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path where the polygon annotations will be saved as a JSON file.",
)
arguments = parser.parse_args()
main(
source_path=arguments.source_path,
zone_configuration_path=arguments.zone_configuration_path,
)
from jsonargparse import auto_cli, set_parsing_settings
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

35
examples/time_in_zone/scripts/stream_from_file.py
View File

@ -1,4 +1,3 @@
import argparse
import os
import subprocess
import tempfile
@ -6,12 +5,20 @@ from glob import glob
from threading import Thread
import yaml
from jsonargparse import auto_cli
SERVER_CONFIG = {"protocols": ["tcp"], "paths": {"all": {"source": "publisher"}}}
BASE_STREAM_URL = "rtsp://localhost:8554/live"
def main(video_directory: str, number_of_streams: int) -> None:
def main(video_directory: str, number_of_streams: int = 6) -> None:
"""
Script to stream videos using RTSP protocol.
Args:
video_directory: Directory containing video files to stream.
number_of_streams: Number of video files to stream.
"""
video_files = find_video_files_in_directory(video_directory, number_of_streams)
try:
with tempfile.TemporaryDirectory() as temporary_directory:
@ -82,23 +89,7 @@ def run_command(command: list) -> int:
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Script to stream videos using RTSP protocol."
)
parser.add_argument(
"--video_directory",
type=str,
required=True,
help="Directory containing video files to stream.",
)
parser.add_argument(
"--number_of_streams",
type=int,
default=6,
help="Number of video files to stream.",
)
arguments = parser.parse_args()
main(
video_directory=arguments.video_directory,
number_of_streams=arguments.number_of_streams,
)
from jsonargparse import auto_cli, set_parsing_settings
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

94
examples/time_in_zone/ultralytics_file_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from ultralytics import YOLO
@ -16,14 +14,26 @@ LABEL_ANNOTATOR = sv.LabelAnnotator(
def main(
source_video_path: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: list[int],
source_video_path: str,
weights: str = "yolov8s.pt",
device: str = "cpu",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
) -> None:
"""
Calculating detections dwell time in zones, using video file.
Args:
zone_configuration_path: Path to the zone configuration JSON file
source_video_path: Path to the source video file
weights: Path to the model weights file
device: Computation device ('cpu', 'mps' or 'cuda')
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
"""
model = YOLO(weights)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
@ -40,10 +50,15 @@ def main(
timers = [FPSBasedTimer(video_info.fps) for _ in zones]
for frame in frames_generator:
results = model(frame, verbose=False, device=device, conf=confidence)[0]
results = model(
frame,
verbose=False,
device=device,
conf=confidence_threshold,
iou=iou_threshold,
)[0]
detections = sv.Detections.from_ultralytics(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = detections.with_nms(threshold=iou)
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
@ -80,60 +95,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using video file."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--source_video_path",
type=str,
required=True,
help="Path to the source video file.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
source_video_path=args.source_video_path,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

94
examples/time_in_zone/ultralytics_naive_stream_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from ultralytics import YOLO
@ -16,14 +14,26 @@ LABEL_ANNOTATOR = sv.LabelAnnotator(
def main(
rtsp_url: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: list[int],
rtsp_url: str,
weights: str = "yolov8s.pt",
device: str = "cpu",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
) -> None:
"""
Calculating detections dwell time in zones, using RTSP stream.
Args:
zone_configuration_path: Path to the zone configuration JSON file
rtsp_url: Complete RTSP URL for the video stream
weights: Path to the model weights file
device: Computation device ('cpu', 'mps' or 'cuda')
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
"""
model = YOLO(weights)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
@ -43,10 +53,15 @@ def main(
fps_monitor.tick()
fps = fps_monitor.fps
results = model(frame, verbose=False, device=device, conf=confidence)[0]
results = model(
frame,
verbose=False,
device=device,
conf=confidence_threshold,
iou=iou_threshold,
)[0]
detections = sv.Detections.from_ultralytics(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = detections.with_nms(threshold=iou)
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
@ -90,60 +105,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)

95
examples/time_in_zone/ultralytics_stream_example.py
View File

@ -1,5 +1,3 @@
import argparse
import cv2
import numpy as np
from inference import InferencePipeline
@ -77,19 +75,37 @@ class CustomSink:
def main(
rtsp_url: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: list[int],
rtsp_url: str,
weights: str = "yolov8s.pt",
device: str = "cpu",
confidence_threshold: float = 0.3,
iou_threshold: float = 0.7,
classes: list[int] = [],
) -> None:
"""
Calculating detections dwell time in zones, using RTSP stream.
Args:
zone_configuration_path: Path to the zone configuration JSON file
rtsp_url: Complete RTSP URL for the video stream
weights: Path to the model weights file
device: Computation device ('cpu', 'mps' or 'cuda')
confidence_threshold: Confidence level for detections (0 to 1)
iou_threshold: IOU threshold for non-max suppression
classes: List of class IDs to track. If empty, all classes are tracked
"""
model = YOLO(weights)
def inference_callback(frame: VideoFrame) -> sv.Detections:
results = model(frame.image, verbose=False, conf=confidence, device=device)[0]
return sv.Detections.from_ultralytics(results).with_nms(threshold=iou)
results = model(
frame.image,
verbose=False,
conf=confidence_threshold,
iou=iou_threshold,
device=device,
)[0]
return sv.Detections.from_ultralytics(results)
sink = CustomSink(zone_configuration_path=zone_configuration_path, classes=classes)
@ -108,60 +124,7 @@ def main(
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
from jsonargparse import auto_cli, set_parsing_settings
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)
set_parsing_settings(parse_optionals_as_positionals=True)
auto_cli(main, as_positional=False)