examples/traffic_analysis: improve CLI argument handling (#2059)

* 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>
This commit is contained in:
Jirka Borovec 2026-02-03 20:56:56 +09:00 committed by GitHub
parent 0efb55ec19
commit 66c49b2e80
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 62 additions and 88 deletions

View File

@ -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)

View File

@ -3,3 +3,4 @@ inference
supervision
tqdm
ultralytics
jsonargparse[signatures]

View File

@ -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)