From 66c49b2e80dc2bca6f51138215071ca97ac90c03 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Tue, 3 Feb 2026 20:56:56 +0900 Subject: [PATCH] examples/traffic_analysis: improve CLI argument handling (#2059) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(traffic_analysis): improve CLI argument handling * fix(pre_commit): 🎨 auto format pre-commit hooks --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../traffic_analysis/inference_example.py | 82 ++++++++----------- examples/traffic_analysis/requirements.txt | 1 + .../traffic_analysis/ultralytics_example.py | 67 +++++++-------- 3 files changed, 62 insertions(+), 88 deletions(-) diff --git a/examples/traffic_analysis/inference_example.py b/examples/traffic_analysis/inference_example.py index f7530ce7..85c1ff4a 100644 --- a/examples/traffic_analysis/inference_example.py +++ b/examples/traffic_analysis/inference_example.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse import os from collections.abc import Iterable @@ -180,62 +179,47 @@ class VideoProcessor: return self.annotate_frame(frame, detections) -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Traffic Flow Analysis with Inference and ByteTrack" - ) +def main( + source_video_path: str, + target_video_path: str, + roboflow_api_key: str, + model_id: str = "vehicle-count-in-drone-video/6", + confidence_threshold: float = 0.3, + iou_threshold: float = 0.7, +) -> None: + """ + Traffic Flow Analysis with Inference and ByteTrack. - parser.add_argument( - "--model_id", - default="vehicle-count-in-drone-video/6", - help="Roboflow model ID", - type=str, - ) - parser.add_argument( - "--roboflow_api_key", - default=None, - help="Roboflow API KEY", - type=str, - ) - parser.add_argument( - "--source_video_path", - required=True, - help="Path to the source video file", - type=str, - ) - parser.add_argument( - "--target_video_path", - default=None, - help="Path to the target video file (output)", - type=str, - ) - parser.add_argument( - "--confidence_threshold", - default=0.3, - help="Confidence threshold for the model", - type=float, - ) - parser.add_argument( - "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float - ) - - args = parser.parse_args() - - api_key = args.roboflow_api_key + Args: + source_video_path: Path to the source video file + target_video_path: Path to the target video file (output) + roboflow_api_key: Roboflow API key + model_id: Roboflow model ID + confidence_threshold: Confidence threshold for the model + iou_threshold: IOU threshold for the model + """ + api_key = roboflow_api_key api_key = os.environ.get("ROBOFLOW_API_KEY", api_key) if api_key is None: raise ValueError( "Roboflow API KEY is missing. Please provide it as an argument or set the " "ROBOFLOW_API_KEY environment variable." ) - args.roboflow_api_key = api_key + roboflow_api_key = api_key processor = VideoProcessor( - roboflow_api_key=args.roboflow_api_key, - model_id=args.model_id, - source_video_path=args.source_video_path, - target_video_path=args.target_video_path, - confidence_threshold=args.confidence_threshold, - iou_threshold=args.iou_threshold, + roboflow_api_key=roboflow_api_key, + model_id=model_id, + source_video_path=source_video_path, + target_video_path=target_video_path, + confidence_threshold=confidence_threshold, + iou_threshold=iou_threshold, ) processor.process_video() + + +if __name__ == "__main__": + from jsonargparse import auto_cli, set_parsing_settings + + set_parsing_settings(parse_optionals_as_positionals=True) + auto_cli(main, as_positional=False) diff --git a/examples/traffic_analysis/requirements.txt b/examples/traffic_analysis/requirements.txt index a8f583e8..f3254183 100644 --- a/examples/traffic_analysis/requirements.txt +++ b/examples/traffic_analysis/requirements.txt @@ -3,3 +3,4 @@ inference supervision tqdm ultralytics +jsonargparse[signatures] diff --git a/examples/traffic_analysis/ultralytics_example.py b/examples/traffic_analysis/ultralytics_example.py index 4df89405..9db82a65 100644 --- a/examples/traffic_analysis/ultralytics_example.py +++ b/examples/traffic_analysis/ultralytics_example.py @@ -1,6 +1,5 @@ from __future__ import annotations -import argparse from collections.abc import Iterable import cv2 @@ -177,45 +176,35 @@ class VideoProcessor: return self.annotate_frame(frame, detections) -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Traffic Flow Analysis with YOLO and ByteTrack" - ) +def main( + source_weights_path: str, + source_video_path: str, + target_video_path: str, + confidence_threshold: float = 0.3, + iou_threshold: float = 0.7, +) -> None: + """ + Traffic Flow Analysis with YOLO and ByteTrack. - parser.add_argument( - "--source_weights_path", - required=True, - help="Path to the source weights file", - type=str, - ) - parser.add_argument( - "--source_video_path", - required=True, - help="Path to the source video file", - type=str, - ) - parser.add_argument( - "--target_video_path", - default=None, - help="Path to the target video file (output)", - type=str, - ) - parser.add_argument( - "--confidence_threshold", - default=0.3, - help="Confidence threshold for the model", - type=float, - ) - parser.add_argument( - "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float - ) - - args = parser.parse_args() + Args: + source_weights_path: Path to the source weights file + source_video_path: Path to the source video file + target_video_path: Path to the target video file (output) + confidence_threshold: Confidence threshold for the model + iou_threshold: IOU threshold for the model + """ processor = VideoProcessor( - source_weights_path=args.source_weights_path, - source_video_path=args.source_video_path, - target_video_path=args.target_video_path, - confidence_threshold=args.confidence_threshold, - iou_threshold=args.iou_threshold, + source_weights_path=source_weights_path, + source_video_path=source_video_path, + target_video_path=target_video_path, + confidence_threshold=confidence_threshold, + iou_threshold=iou_threshold, ) processor.process_video() + + +if __name__ == "__main__": + from jsonargparse import auto_cli, set_parsing_settings + + set_parsing_settings(parse_optionals_as_positionals=True) + auto_cli(main, as_positional=False)