From 6a1b2b8eb4823f961464edf1a544bfa6a693e112 Mon Sep 17 00:00:00 2001 From: Jirka Borovec <6035284+Borda@users.noreply.github.com> Date: Tue, 3 Feb 2026 19:35:09 +0900 Subject: [PATCH] examples/count_people_in_zone: resolve IoU threshold & improve CLI (#2054) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(count_people_in_zone): resolve IoU threshold & improve CLI * fix(pre_commit): 🎨 auto format pre-commit hooks * Apply suggestions from code review --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../count_people_in_zone/inference_example.py | 107 +++++++----------- .../count_people_in_zone/requirements.txt | 1 + .../ultralytics_example.py | 95 +++++++--------- 3 files changed, 86 insertions(+), 117 deletions(-) diff --git a/examples/count_people_in_zone/inference_example.py b/examples/count_people_in_zone/inference_example.py index 2225594f..677af693 100644 --- a/examples/count_people_in_zone/inference_example.py +++ b/examples/count_people_in_zone/inference_example.py @@ -1,4 +1,3 @@ -import argparse import json import os @@ -62,7 +61,10 @@ def initiate_annotators( def detect( - frame: np.ndarray, model: RoboflowInferenceModel, confidence_threshold: float = 0.5 + frame: np.ndarray, + model: RoboflowInferenceModel, + confidence_threshold: float = 0.5, + iou_threshold: float = 0.7, ) -> sv.Detections: """ Detect objects in a frame using Inference model, filtering detections by class ID @@ -73,7 +75,8 @@ def detect( model (RoboflowInferenceModel): The Inference model used for processing the frame. confidence_threshold (float): The confidence threshold for filtering - detections. Default is 0.5. + detections. + iou_threshold (float): The IoU threshold for non-maximum suppression. Returns: sv.Detections: Filtered detections after processing the frame with the Inference @@ -83,7 +86,7 @@ def detect( This function is specifically tailored for an Inference model and assumes class ID 0 for filtering. """ - results = model.infer(frame)[0] + results = model.infer(frame, confidence=confidence_threshold, iou=iou_threshold)[0] detections = sv.Detections.from_inference(results) filter_by_class = detections.class_id == 0 filter_by_confidence = detections.confidence > confidence_threshold @@ -124,78 +127,49 @@ def annotate( return annotated_frame -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Counting people in zones with Inference and Supervision" - ) +def main( + zone_configuration_path: str, + source_video_path: str, + model_id: str = "yolov8x-1280", + roboflow_api_key: str | None = None, + target_video_path: str | None = None, + confidence_threshold: float = 0.3, + iou_threshold: float = 0.7, +): + """ + Counting people in zones with Inference and Supervision. - parser.add_argument( - "--zone_configuration_path", - required=True, - help="Path to the zone configuration JSON file", - type=str, - ) - parser.add_argument( - "--model_id", - default="yolov8x-1280", - 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: + zone_configuration_path: Path to the zone configuration JSON file + source_video_path: Path to the source video file + model_id: Roboflow model ID + roboflow_api_key: Roboflow API KEY + target_video_path: Path to the target video file (output) + 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 - video_info = sv.VideoInfo.from_video_path(args.source_video_path) - polygons = load_zones_config(args.zone_configuration_path) + video_info = sv.VideoInfo.from_video_path(source_video_path) + polygons = load_zones_config(zone_configuration_path) zones, zone_annotators, box_annotators = initiate_annotators( polygons=polygons, resolution_wh=video_info.resolution_wh ) - model = get_roboflow_model(model_id=args.model_id, api_key=args.roboflow_api_key) + model = get_roboflow_model(model_id=model_id, api_key=roboflow_api_key) - frames_generator = sv.get_video_frames_generator(args.source_video_path) - if args.target_video_path is not None: - with sv.VideoSink(args.target_video_path, video_info) as sink: + frames_generator = sv.get_video_frames_generator(source_video_path) + if target_video_path is not None: + with sv.VideoSink(target_video_path, video_info) as sink: for frame in tqdm(frames_generator, total=video_info.total_frames): - detections = detect(frame, model, args.confidence_threshold) + detections = detect(frame, model, confidence_threshold, iou_threshold) annotated_frame = annotate( frame=frame, zones=zones, @@ -206,7 +180,7 @@ if __name__ == "__main__": sink.write_frame(annotated_frame) else: for frame in tqdm(frames_generator, total=video_info.total_frames): - detections = detect(frame, model, args.confidence_threshold) + detections = detect(frame, model, confidence_threshold, iou_threshold) annotated_frame = annotate( frame=frame, zones=zones, @@ -219,3 +193,10 @@ if __name__ == "__main__": break cv2.destroyAllWindows() + + +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/count_people_in_zone/requirements.txt b/examples/count_people_in_zone/requirements.txt index a8f583e8..f3254183 100644 --- a/examples/count_people_in_zone/requirements.txt +++ b/examples/count_people_in_zone/requirements.txt @@ -3,3 +3,4 @@ inference supervision tqdm ultralytics +jsonargparse[signatures] diff --git a/examples/count_people_in_zone/ultralytics_example.py b/examples/count_people_in_zone/ultralytics_example.py index 22909480..35a09616 100644 --- a/examples/count_people_in_zone/ultralytics_example.py +++ b/examples/count_people_in_zone/ultralytics_example.py @@ -1,4 +1,3 @@ -import argparse import json import cv2 @@ -60,7 +59,10 @@ def initiate_annotators( def detect( - frame: np.ndarray, model: YOLO, confidence_threshold: float = 0.5 + frame: np.ndarray, + model: YOLO, + confidence_threshold: float = 0.5, + iou_threshold: float = 0.7, ) -> sv.Detections: """ Detect objects in a frame using a YOLO model, filtering detections by class ID and @@ -71,6 +73,7 @@ def detect( model (YOLO): The YOLO model used for processing the frame. confidence_threshold (float): The confidence threshold for filtering detections. Default is 0.5. + iou_threshold (float): The IoU threshold for non-maximum suppression. Returns: sv.Detections: Filtered detections after processing the frame with the YOLO @@ -80,7 +83,9 @@ def detect( This function is specifically tailored for a YOLO model and assumes class ID 0 for filtering. """ - results = model(frame, imgsz=1280, verbose=False)[0] + results = model( + frame, conf=confidence_threshold, iou=iou_threshold, imgsz=1280, verbose=False + )[0] detections = sv.Detections.from_ultralytics(results) filter_by_class = detections.class_id == 0 filter_by_confidence = detections.confidence > confidence_threshold @@ -121,63 +126,38 @@ def annotate( return annotated_frame -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Counting people in zones with YOLO and Supervision" - ) +def main( + zone_configuration_path: str, + source_video_path: str, + source_weights_path: str = "yolo11x.pt", + target_video_path: str | None = None, + confidence_threshold: float = 0.3, + iou_threshold: float = 0.7, +): + """ + Counting people in zones with YOLO and Supervision. - parser.add_argument( - "--zone_configuration_path", - required=True, - help="Path to the zone configuration JSON file", - type=str, - ) - parser.add_argument( - "--source_weights_path", - default="yolo11x.pt", - 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() - - video_info = sv.VideoInfo.from_video_path(args.source_video_path) - polygons = load_zones_config(args.zone_configuration_path) + Args: + zone_configuration_path: Path to the zone configuration JSON file + source_video_path: Path to the source video file + source_weights_path: Path to the source weights 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 + """ + video_info = sv.VideoInfo.from_video_path(source_video_path) + polygons = load_zones_config(zone_configuration_path) zones, zone_annotators, box_annotators = initiate_annotators( polygons=polygons, resolution_wh=video_info.resolution_wh ) - model = YOLO(args.source_weights_path) + model = YOLO(source_weights_path) - frames_generator = sv.get_video_frames_generator(args.source_video_path) - if args.target_video_path is not None: - with sv.VideoSink(args.target_video_path, video_info) as sink: + frames_generator = sv.get_video_frames_generator(source_video_path) + if target_video_path is not None: + with sv.VideoSink(target_video_path, video_info) as sink: for frame in tqdm(frames_generator, total=video_info.total_frames): - detections = detect(frame, model, args.confidence_threshold) + detections = detect(frame, model, confidence_threshold, iou_threshold) annotated_frame = annotate( frame=frame, zones=zones, @@ -188,7 +168,7 @@ if __name__ == "__main__": sink.write_frame(annotated_frame) else: for frame in tqdm(frames_generator, total=video_info.total_frames): - detections = detect(frame, model, args.confidence_threshold) + detections = detect(frame, model, confidence_threshold, iou_threshold) annotated_frame = annotate( frame=frame, zones=zones, @@ -201,3 +181,10 @@ if __name__ == "__main__": break cv2.destroyAllWindows() + + +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)