Annotator types and docs: MyPy, consistent 'optional', small tidy-ups

docs/detection/annotators.md

@@ -384,7 +384,7 @@ status: new
     trace_annotator = sv.TraceAnnotator()
 
     video_info = sv.VideoInfo.from_video_path(video_path='...')
-    frames_generator = get_video_frames_generator(source_path='...')
+    frames_generator = sv.get_video_frames_generator(source_path='...')
     tracker = sv.ByteTrack()
 
     with sv.VideoSink(target_path='...', video_info=video_info) as sink:
@@ -415,7 +415,7 @@ status: new
     heat_map_annotator = sv.HeatMapAnnotator()
 
     video_info = sv.VideoInfo.from_video_path(video_path='...')
-    frames_generator = get_video_frames_generator(source_path='...')
+    frames_generator = sv.get_video_frames_generator(source_path='...')
 
     with sv.VideoSink(target_path='...', video_info=video_info) as sink:
         for frame in frames_generator:

examples/count_people_in_zone/inference_example.py

@@ -75,7 +75,7 @@ def detect(
         frame (np.ndarray): The frame to process, expected to be a NumPy array.
         model (RoboflowInferenceModel): The Inference model used for processing the
             frame.
-        confidence_threshold (float, optional): The confidence threshold for filtering
+        confidence_threshold (float): The confidence threshold for filtering
             detections. Default is 0.5.
 
     Returns:

examples/count_people_in_zone/ultralytics_example.py

@@ -72,7 +72,7 @@ def detect(
     Args:
         frame (np.ndarray): The frame to process, expected to be a NumPy array.
         model (YOLO): The YOLO model used for processing the frame.
-        confidence_threshold (float, optional): The confidence threshold for filtering
+        confidence_threshold (float): The confidence threshold for filtering
             detections. Default is 0.5.
 
     Returns:

examples/time_in_zone/utils/timers.py

@@ -22,7 +22,7 @@ def __init__(self, fps: int = 30) -> None:
         """Initializes the FPSBasedTimer with the specified frames per second rate.
 
         Args:
-            fps (int, optional): The frame rate of the video stream. Defaults to 30.
+            fps (int): The frame rate of the video stream. Defaults to 30.
         """
         self.fps = fps
         self.frame_id = 0

supervision/annotators/core.py

@@ -3,7 +3,8 @@
 
 import cv2
 import numpy as np
-from PIL import ImageDraw, ImageFont
+import numpy.typing as npt
+from PIL import Image, ImageDraw, ImageFont
 
 from supervision.annotators.base import BaseAnnotator, ImageType
 from supervision.annotators.utils import (
@@ -87,6 +88,7 @@ def annotate(
         ![bounding-box-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/bounding-box-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             color = resolve_color(
@@ -172,6 +174,7 @@ def annotate(
         ![bounding-box-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/bounding-box-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             color = resolve_color(
@@ -256,12 +259,13 @@ def annotate(
             )
             ```
         """  # noqa E501 // docs
-
+        assert isinstance(scene, np.ndarray)
         if detections.data is None or ORIENTED_BOX_COORDINATES not in detections.data:
             return scene
+        obb_boxes = np.array(detections.data[ORIENTED_BOX_COORDINATES]).astype(int)
 
         for detection_idx in range(len(detections)):
-            bbox = np.intp(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx])
+            obb = obb_boxes[detection_idx]
             color = resolve_color(
                 color=self.color,
                 detections=detections,
@@ -271,7 +275,7 @@ def annotate(
                 else custom_color_lookup,
             )
 
-            cv2.drawContours(scene, [bbox], 0, color.as_bgr(), self.thickness)
+            cv2.drawContours(scene, [obb], 0, color.as_bgr(), self.thickness)
 
         return scene
 
@@ -342,6 +346,7 @@ def annotate(
         ![mask-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/mask-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         if detections.mask is None:
             return scene
 
@@ -431,6 +436,7 @@ def annotate(
         ![polygon-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/polygon-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         if detections.mask is None:
             return scene
 
@@ -517,6 +523,7 @@ def annotate(
         ![box-mask-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/box-mask-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         scene_with_boxes = scene.copy()
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
@@ -612,6 +619,7 @@ def annotate(
         ![halo-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/halo-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         if detections.mask is None:
             return scene
         colored_mask = np.zeros_like(scene, dtype=np.uint8)
@@ -711,6 +719,7 @@ def annotate(
         ![ellipse-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/ellipse-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             color = resolve_color(
@@ -802,6 +811,7 @@ def annotate(
         ![box-corner-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/box-corner-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             color = resolve_color(
@@ -891,6 +901,7 @@ def annotate(
         ![circle-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/circle-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             center = ((x1 + x2) // 2, (y1 + y2) // 2)
@@ -983,6 +994,7 @@ def annotate(
         ![dot-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/dot-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         xy = detections.get_anchors_coordinates(anchor=self.position)
         for detection_idx in range(len(detections)):
             color = resolve_color(
@@ -1092,6 +1104,7 @@ def annotate(
         ![label-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/label-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         font = cv2.FONT_HERSHEY_SIMPLEX
         anchors_coordinates = detections.get_anchors_coordinates(
             anchor=self.text_anchor
@@ -1128,8 +1141,8 @@ def annotate(
 
             if labels is not None:
                 text = labels[detection_idx]
-            elif detections[CLASS_NAME_DATA_FIELD] is not None:
-                text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
+            elif CLASS_NAME_DATA_FIELD in detections.data:
+                text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
             elif detections.class_id is not None:
                 text = str(detections.class_id[detection_idx])
             else:
@@ -1308,6 +1321,7 @@ def annotate(
             ```
 
         """
+        assert isinstance(scene, Image.Image)
         draw = ImageDraw.Draw(scene)
         anchors_coordinates = detections.get_anchors_coordinates(
             anchor=self.text_anchor
@@ -1343,8 +1357,8 @@ def annotate(
 
             if labels is not None:
                 text = labels[detection_idx]
-            elif detections[CLASS_NAME_DATA_FIELD] is not None:
-                text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
+            elif CLASS_NAME_DATA_FIELD in detections.data:
+                text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
             elif detections.class_id is not None:
                 text = str(detections.class_id[detection_idx])
             else:
@@ -1440,6 +1454,7 @@ def annotate(
         ![blur-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/blur-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         image_height, image_width = scene.shape[:2]
         clipped_xyxy = clip_boxes(
             xyxy=detections.xyxy, resolution_wh=(image_width, image_height)
@@ -1538,8 +1553,14 @@ def annotate(
         ![trace-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/trace-annotator-example-purple.png)
         """
-        self.trace.put(detections)
+        assert isinstance(scene, np.ndarray)
+        if detections.tracker_id is None:
+            raise ValueError(
+                "The `tracker_id` field is missing in the provided detections."
+                " See more: https://supervision.roboflow.com/latest/how_to/track_objects"
+            )
 
+        self.trace.put(detections)
         for detection_idx in range(len(detections)):
             tracker_id = int(detections.tracker_id[detection_idx])
             color = resolve_color(
@@ -1592,9 +1613,9 @@ def __init__(
         self.opacity = opacity
         self.radius = radius
         self.kernel_size = kernel_size
-        self.heat_mask = None
         self.top_hue = top_hue
         self.low_hue = low_hue
+        self.heat_mask: Optional[npt.NDArray[np.float32]] = None
 
     @ensure_cv2_image_for_annotation
     def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
@@ -1621,7 +1642,7 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
             heat_map_annotator = sv.HeatMapAnnotator()
 
             video_info = sv.VideoInfo.from_video_path(video_path='...')
-            frames_generator = get_video_frames_generator(source_path='...')
+            frames_generator = sv.get_video_frames_generator(source_path='...')
 
             with sv.VideoSink(target_path='...', video_info=video_info) as sink:
                for frame in frames_generator:
@@ -1636,12 +1657,20 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
         ![heatmap-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/heat-map-annotator-example-purple.png)
         """
-
+        assert isinstance(scene, np.ndarray)
         if self.heat_mask is None:
-            self.heat_mask = np.zeros(scene.shape[:2])
+            self.heat_mask = np.zeros(scene.shape[:2], dtype=np.float32)
+
         mask = np.zeros(scene.shape[:2])
         for xy in detections.get_anchors_coordinates(self.position):
-            cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1)
+            x, y = int(xy[0]), int(xy[1])
+            cv2.circle(
+                img=mask,
+                center=(x, y),
+                radius=self.radius,
+                color=(1,),
+                thickness=-1,  # fill
+            )
         self.heat_mask = mask + self.heat_mask
         temp = self.heat_mask.copy()
         temp = self.low_hue - temp / temp.max() * (self.low_hue - self.top_hue)
@@ -1709,6 +1738,7 @@ def annotate(
         ![pixelate-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/pixelate-annotator-example-10.png)
         """
+        assert isinstance(scene, np.ndarray)
         image_height, image_width = scene.shape[:2]
         clipped_xyxy = clip_boxes(
             xyxy=detections.xyxy, resolution_wh=(image_width, image_height)
@@ -1802,6 +1832,7 @@ def annotate(
         ![triangle-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/triangle-annotator-example.png)
         """
+        assert isinstance(scene, np.ndarray)
         xy = detections.get_anchors_coordinates(anchor=self.position)
         for detection_idx in range(len(detections)):
             color = resolve_color(
@@ -1902,7 +1933,7 @@ def annotate(
         ![round-box-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/round-box-annotator-example-purple.png)
         """
-
+        assert isinstance(scene, np.ndarray)
         for detection_idx in range(len(detections)):
             x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
             color = resolve_color(
@@ -1975,7 +2006,7 @@ def __init__(
         border_color: Color = Color.BLACK,
         position: Position = Position.TOP_CENTER,
         color_lookup: ColorLookup = ColorLookup.CLASS,
-        border_thickness: int = None,
+        border_thickness: Optional[int] = None,
     ):
         """
         Args:
@@ -1987,7 +2018,7 @@ def __init__(
             position (Position): The anchor position of drawing the percentage bar.
             color_lookup (ColorLookup): Strategy for mapping colors to annotations.
                 Options are `INDEX`, `CLASS`, `TRACK`.
-            border_thickness (int): The thickness of the border lines.
+            border_thickness (Optional[int]): The thickness of the border lines.
         """
         self.height: int = height
         self.width: int = width
@@ -2046,9 +2077,9 @@ def annotate(
         ![percentage-bar-example](https://media.roboflow.com/
         supervision-annotator-examples/percentage-bar-annotator-example-purple.png)
         """
-        self.validate_custom_values(
-            custom_values=custom_values, detections_count=len(detections)
-        )
+        assert isinstance(scene, np.ndarray)
+        self.validate_custom_values(custom_values=custom_values, detections=detections)
+
         anchors = detections.get_anchors_coordinates(anchor=self.position)
         for detection_idx in range(len(detections)):
             anchor = anchors[detection_idx]
@@ -2059,11 +2090,11 @@ def annotate(
             )
             border_width = border_coordinates[1][0] - border_coordinates[0][0]
 
-            value = (
-                custom_values[detection_idx]
-                if custom_values is not None
-                else detections.confidence[detection_idx]
-            )
+            if custom_values is not None:
+                value = custom_values[detection_idx]
+            else:
+                assert detections.confidence is not None  # MyPy type hint
+                value = detections.confidence[detection_idx]
 
             color = resolve_color(
                 color=self.color,
@@ -2123,15 +2154,23 @@ def calculate_border_coordinates(
 
     @staticmethod
     def validate_custom_values(
-        custom_values: Optional[Union[np.ndarray, List[float]]], detections_count: int
+        custom_values: Optional[Union[np.ndarray, List[float]]], detections: Detections
     ) -> None:
-        if custom_values is not None:
+        if custom_values is None:
+            if detections.confidence is None:
+                raise ValueError(
+                    "The provided detections do not contain confidence values. "
+                    "Please provide `custom_values` or ensure that the detections "
+                    "contain confidence values (e.g. by using a different model)."
+                )
+
+        else:
             if not isinstance(custom_values, (np.ndarray, list)):
                 raise TypeError(
                     "custom_values must be either a numpy array or a list of floats."
                 )
 
-            if len(custom_values) != detections_count:
+            if len(custom_values) != len(detections):
                 raise ValueError(
                     "The length of custom_values must match the number of detections."
                 )
@@ -2211,6 +2250,7 @@ def annotate(
             )
             ```
         """
+        assert isinstance(scene, np.ndarray)
         crops = [
             crop_image(image=scene, xyxy=xyxy) for xyxy in detections.xyxy.astype(int)
         ]
@@ -2349,6 +2389,7 @@ def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
         ![background-overlay-annotator-example](https://media.roboflow.com/
         supervision-annotator-examples/background-color-annotator-example-purple.png)
         """
+        assert isinstance(scene, np.ndarray)
         colored_mask = np.full_like(scene, self.color.as_bgr(), dtype=np.uint8)
 
         cv2.addWeighted(