examples/count_people_in_zone: resolve IoU threshold & improve CLI (#2054)

* 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>

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)

examples/count_people_in_zone/requirements.txt

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

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)