👊 bump version before release

fix/041_index_error_on_polygon_zones

docs/detection_utils.md

@@ -1,3 +1,11 @@
 ## generate_2d_mask
 
-:::supervision.detection.utils.generate_2d_mask
+:::supervision.detection.utils.generate_2d_mask
+
+## box_iou_batch
+
+:::supervision.detection.utils.box_iou_batch
+
+## non_max_suppression
+
+:::supervision.detection.utils.non_max_suppression

docs/draw.md

@@ -1,7 +0,0 @@
-## draw_line
-
-:::supervision.draw.utils.draw_line
-
-## draw_rectangle
-
-:::supervision.draw.utils.draw_rectangle

docs/draw_utils.md

@@ -0,0 +1,19 @@
+## draw_line
+
+:::supervision.draw.utils.draw_line
+
+## draw_rectangle
+
+:::supervision.draw.utils.draw_rectangle
+
+## draw_filled_rectangle
+
+:::supervision.draw.utils.draw_filled_rectangle
+
+## draw_polygon
+
+:::supervision.draw.utils.draw_polygon
+
+## draw_text
+
+:::supervision.draw.utils.draw_text

mkdocs.yml

@@ -28,8 +28,8 @@ nav:
     - Detection:
         - Core: detection_core.md
         - Utils: detection_utils.md
-    - Draw: draw.md
-    - Geometry: geometry.md
+    - Draw:
+        - Utils: draw_utils.md
     - Notebook: notebook.md
 
 theme:

supervision/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 
 from supervision.detection.core import BoxAnnotator, Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator

supervision/detection/core.py

@@ -93,14 +93,15 @@ class Detections:
         )
 
     @classmethod
-    def from_yolov5(cls, yolov5_detections):
+    def from_yolov5(cls, yolov5_results) -> Detections:
         """
         Creates a Detections instance from a YOLOv5 output Detections
 
-        Attributes:
-            yolov5_detections (yolov5.models.common.Detections): The output Detections instance from YOLOv5
+        Args:
+            yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
 
         Returns:
+            Detections: A new Detections object.
 
         Example:
             ```python
@@ -112,7 +113,7 @@ class Detections:
             >>> detections = Detections.from_yolov5(results)
             ```
         """
-        yolov5_detections_predictions = yolov5_detections.pred[0].cpu().cpu().numpy()
+        yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
         return cls(
             xyxy=yolov5_detections_predictions[:, :4],
             confidence=yolov5_detections_predictions[:, 4],
@@ -120,14 +121,15 @@ class Detections:
         )
 
     @classmethod
-    def from_yolov8(cls, yolov8_results):
+    def from_yolov8(cls, yolov8_results) -> Detections:
         """
         Creates a Detections instance from a YOLOv8 output Results
 
-        Attributes:
+        Args:
             yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
 
         Returns:
+            Detections: A new Detections object.
 
         Example:
             ```python
@@ -146,7 +148,13 @@ class Detections:
         )
 
     @classmethod
-    def from_transformers(cls, transformers_results: dict):
+    def from_transformers(cls, transformers_results: dict) -> Detections:
+        """
+        Creates a Detections instance from Object Detection Transformer output Results
+
+        Returns:
+            Detections: A new Detections object.
+        """
         return cls(
             xyxy=transformers_results["boxes"].cpu().numpy(),
             confidence=transformers_results["scores"].cpu().numpy(),
@@ -154,7 +162,7 @@ class Detections:
         )
 
     @classmethod
-    def from_detectron2(cls, detectron2_results):
+    def from_detectron2(cls, detectron2_results) -> Detections:
         return cls(
             xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
             confidence=detectron2_results["instances"].scores.cpu().numpy(),
@@ -165,7 +173,7 @@ class Detections:
         )
 
     @classmethod
-    def from_coco_annotations(cls, coco_annotation: dict):
+    def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
         xyxy, class_id = [], []
 
         for annotation in coco_annotation:
@@ -175,40 +183,11 @@ class Detections:
 
         return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
 
-    def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]:
-        """
-        Filter the detections by applying a mask.
-
-        Attributes:
-            mask (np.ndarray): A mask of shape `(n,)` containing a boolean value for each detection indicating if it should be included in the filtered detections
-            inplace (bool): If True, the original data will be modified and self will be returned.
-
-        Returns:
-            Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise.
-        """
-        if inplace:
-            self.xyxy = self.xyxy[mask]
-            self.confidence = self.confidence[mask]
-            self.class_id = self.class_id[mask]
-            self.tracker_id = (
-                self.tracker_id[mask] if self.tracker_id is not None else None
-            )
-            return self
-        else:
-            return Detections(
-                xyxy=self.xyxy[mask],
-                confidence=self.confidence[mask],
-                class_id=self.class_id[mask],
-                tracker_id=self.tracker_id[mask]
-                if self.tracker_id is not None
-                else None,
-            )
-
     def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
         """
         Returns the bounding box coordinates for a specific anchor.
 
-        Properties:
+        Args:
             anchor (Position): Position of bounding box anchor for which to return the coordinates.
 
         Returns:
@@ -246,13 +225,53 @@ class Detections:
 
     @property
     def area(self) -> np.ndarray:
+        """
+        Calculate the area of each bounding box in the set of object detections.
+
+        Returns:
+            np.ndarray: An array of floats containing the area of each bounding box in the format of (area_1, area_2, ..., area_n), where n is the number of detections.
+        """
         return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0])
 
-    def with_nms(self, threshold: float = 0.5) -> Detections:
+    def with_nms(
+        self, threshold: float = 0.5, class_agnostic: bool = False
+    ) -> Detections:
+        """
+        Perform non-maximum suppression on the current set of object detections.
+
+        Args:
+            threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5.
+            class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False.
+
+        Returns:
+            Detections: A new Detections object containing the subset of detections after non-maximum suppression.
+
+        Raises:
+            AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False.
+        """
+        if len(self) == 0:
+            return self
+
         assert (
             self.confidence is not None
         ), f"Detections confidence must be given for NMS to be executed."
-        indices = non_max_suppression(self.xyxy, self.confidence, threshold=threshold)
+
+        if class_agnostic:
+            predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
+            indices = non_max_suppression(
+                predictions=predictions, iou_threshold=threshold
+            )
+            return self[indices]
+
+        assert self.class_id is not None, (
+            f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic "
+            f"NMS set class_agnostic=True."
+        )
+
+        predictions = np.hstack(
+            (self.xyxy, self.confidence.reshape(-1, 1), self.class_id.reshape(-1, 1))
+        )
+        indices = non_max_suppression(predictions=predictions, iou_threshold=threshold)
         return self[indices]
 
 

supervision/detection/polygon_zone.py

@@ -1,10 +1,11 @@
+from dataclasses import replace
 from typing import Optional, Tuple
 
 import cv2
 import numpy as np
 
 from supervision import Detections
-from supervision.detection.utils import generate_2d_mask
+from supervision.detection.utils import clip_boxes, generate_2d_mask
 from supervision.draw.color import Color
 from supervision.draw.utils import draw_polygon, draw_text
 from supervision.geometry.core import Position
@@ -21,17 +22,22 @@ class PolygonZone:
         self.polygon = polygon
         self.frame_resolution_wh = frame_resolution_wh
         self.triggering_position = triggering_position
-        self.mask = generate_2d_mask(polygon=polygon, resolution_wh=frame_resolution_wh)
         self.current_count = 0
 
-    def trigger(self, detections: Detections) -> np.ndarray:
-        anchors = (
-            np.ceil(
-                detections.get_anchor_coordinates(anchor=self.triggering_position)
-            ).astype(int)
-            - 1
+        width, height = frame_resolution_wh
+        self.mask = generate_2d_mask(
+            polygon=polygon, resolution_wh=(width + 1, height + 1)
         )
-        is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]]
+
+    def trigger(self, detections: Detections) -> np.ndarray:
+        clipped_xyxy = clip_boxes(
+            boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh
+        )
+        clipped_detections = replace(detections, xyxy=clipped_xyxy)
+        clipped_anchors = np.ceil(
+            clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
+        ).astype(int)
+        is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
         self.current_count = np.sum(is_in_zone)
         return is_in_zone.astype(bool)
 

supervision/detection/utils.py

@@ -7,7 +7,7 @@ import numpy as np
 def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
     """Generate a 2D mask from a polygon.
 
-    Properties:
+    Args:
         polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
         resolution_wh (Tuple[int, int]): The width and height of the desired resolution.
 
@@ -20,38 +20,98 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.
     return mask
 
 
-def non_max_suppression(boxes: np.ndarray, scores: np.ndarray, threshold: float):
-    assert boxes.shape[0] == scores.shape[0]
-    ys1 = boxes[:, 0]
-    xs1 = boxes[:, 1]
-    ys2 = boxes[:, 2]
-    xs2 = boxes[:, 3]
+def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
+    """
+    Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets
+    of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
 
-    areas = (ys2 - ys1) * (xs2 - xs1)
-    scores_indexes = scores.argsort().tolist()
-    boxes_keep_index = []
-    while len(scores_indexes):
-        index = scores_indexes.pop()
-        boxes_keep_index.append(index)
-        if not len(scores_indexes):
-            break
-        iou = compute_iou(
-            boxes[index], boxes[scores_indexes], areas[index], areas[scores_indexes]
-        )
-        filtered_indexes = set((iou > threshold).nonzero()[0])
-        scores_indexes = [
-            v for (i, v) in enumerate(scores_indexes) if i not in filtered_indexes
-        ]
-    return np.array(boxes_keep_index)
+    Args:
+        boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
+        boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.
+
+    Returns:
+        np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects.
+    """
+
+    def box_area(box):
+        return (box[2] - box[0]) * (box[3] - box[1])
+
+    area_true = box_area(boxes_true.T)
+    area_detection = box_area(boxes_detection.T)
+
+    top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
+    bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])
+
+    area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
+    return area_inter / (area_true[:, None] + area_detection - area_inter)
 
 
-def compute_iou(box, boxes, box_area, boxes_area):
-    assert boxes.shape[0] == boxes_area.shape[0]
-    ys1 = np.maximum(box[0], boxes[:, 0])
-    xs1 = np.maximum(box[1], boxes[:, 1])
-    ys2 = np.minimum(box[2], boxes[:, 2])
-    xs2 = np.minimum(box[3], boxes[:, 3])
-    intersections = np.maximum(ys2 - ys1, 0) * np.maximum(xs2 - xs1, 0)
-    unions = box_area + boxes_area - intersections
-    iou = intersections / unions
-    return iou
+def non_max_suppression(
+    predictions: np.ndarray, iou_threshold: float = 0.5
+) -> np.ndarray:
+    """
+    Perform Non-Maximum Suppression (NMS) on object detection predictions.
+
+    Args:
+        predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`.
+        iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression.
+
+    Returns:
+        np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.
+
+    Raises:
+        AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
+    """
+    assert 0 <= iou_threshold <= 1, (
+        f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
+        f"{iou_threshold} given."
+    )
+    rows, columns = predictions.shape
+
+    # add column #5 - category filled with zeros for agnostic nms
+    if columns == 5:
+        predictions = np.c_[predictions, np.zeros(rows)]
+
+    # sort predictions column #4 - score
+    sort_index = np.flip(predictions[:, 4].argsort())
+    predictions = predictions[sort_index]
+
+    boxes = predictions[:, :4]
+    categories = predictions[:, 5]
+    ious = box_iou_batch(boxes, boxes)
+    ious = ious - np.eye(rows)
+
+    keep = np.ones(rows, dtype=bool)
+
+    for index, (iou, category) in enumerate(zip(ious, categories)):
+        if not keep[index]:
+            continue
+
+        # drop detections with iou > iou_threshold and same category as current detections
+        condition = (iou > iou_threshold) & (categories == category)
+        keep = keep & ~condition
+
+    return keep[sort_index.argsort()]
+
+
+def clip_boxes(
+    boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+    """
+    Clips bounding boxes coordinates to fit within the frame resolution.
+
+    Args:
+        boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in
+        the format `(x_min, y_min, x_max, y_max)`.
+        frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the
+        frame.
+
+    Returns:
+        np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates
+        clipped to fit within the frame resolution.
+    """
+    result = np.copy(boxes_xyxy)
+    width, height = frame_resolution_wh
+    result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
+    result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
+    return result

supervision/draw/utils.py

@@ -39,7 +39,7 @@ def draw_rectangle(
     """
     Draws a rectangle on an image.
 
-    Attributes:
+    Parameters:
         scene (np.ndarray): The scene on which the rectangle will be drawn
         rect (Rect): The rectangle to be drawn
         color (Color): The color of the rectangle
@@ -60,14 +60,9 @@ def draw_rectangle(
 
 def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.ndarray:
     """
-    Draws a filled rectangle on the given scene.
+    Draws a filled rectangle on an image.
 
-    :param scene: np.ndarray : The scene on which to draw the rectangle.
-    :param rect: Rect : The rectangle to be drawn.
-    :param color: Color : The color of the rectangle.
-    :return: np.ndarray : The updated scene with the filled rectangle drawn on it.
-
-    Attributes:
+    Parameters:
         scene (np.ndarray): The scene on which the rectangle will be drawn
         rect (Rect): The rectangle to be drawn
         color (Color): The color of the rectangle
@@ -90,7 +85,7 @@ def draw_polygon(
 ) -> np.ndarray:
     """Draw a polygon on a scene.
 
-    Attributes:
+    Parameters:
         scene (np.ndarray): The scene to draw the polygon on.
         polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
         color (Color): The color of the polygon.
@@ -117,9 +112,7 @@ def draw_text(
     background_color: Optional[Color] = None,
 ) -> np.ndarray:
     """
-    Draw text on a scene.
-
-    This function takes in a 2-dimensional numpy ndarray representing an image or scene, and draws text on it using OpenCV's putText function. The text is anchored at a specified Point, and its appearance can be customized using arguments such as color, scale, and font. An optional background color and padding can be specified to draw a rectangle behind the text.
+    Draw text with background on a scene.
 
     Parameters:
         scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene.

test/detection/test_utils.py

@@ -0,0 +1,188 @@
+from contextlib import ExitStack as DoesNotRaise
+from typing import Optional, Tuple
+
+import pytest
+
+import numpy as np
+
+from supervision.detection.utils import non_max_suppression, clip_boxes
+
+
+@pytest.mark.parametrize(
+    "predictions, iou_threshold, expected_result, exception",
+    [
+        (
+            np.empty(shape=(0, 5)),
+            0.5,
+            np.array([]),
+            DoesNotRaise()
+        ),  # single box with no category
+        (
+            np.array([
+                [10.0, 10.0, 40.0, 40.0, 0.8]
+            ]),
+            0.5,
+            np.array([
+                True
+            ]),
+            DoesNotRaise()
+        ),  # single box with no category
+        (
+            np.array([
+                [10.0, 10.0, 40.0, 40.0, 0.8, 0]
+            ]),
+            0.5,
+            np.array([
+                True
+            ]),
+            DoesNotRaise()
+        ),  # single box with category
+        (
+            np.array([
+                [10.0, 10.0, 40.0, 40.0, 0.8],
+                [15.0, 15.0, 40.0, 40.0, 0.9],
+            ]),
+            0.5,
+            np.array([
+                False,
+                True
+            ]),
+            DoesNotRaise()
+        ),  # two boxes with no category
+        (
+            np.array([
+                [10.0, 10.0, 40.0, 40.0, 0.8, 0],
+                [15.0, 15.0, 40.0, 40.0, 0.9, 1],
+            ]),
+            0.5,
+            np.array([
+                True,
+                True
+            ]),
+            DoesNotRaise()
+        ),  # two boxes with different category
+        (
+            np.array([
+                [10.0, 10.0, 40.0, 40.0, 0.8, 0],
+                [15.0, 15.0, 40.0, 40.0, 0.9, 0],
+            ]),
+            0.5,
+            np.array([
+                False,
+                True
+            ]),
+            DoesNotRaise()
+        ),  # two boxes with same category
+        (
+            np.array([
+                [0.0, 0.0, 30.0, 40.0, 0.8],
+                [5.0, 5.0, 35.0, 45.0, 0.9],
+                [10.0, 10.0, 40.0, 50.0, 0.85],
+            ]),
+            0.5,
+            np.array([
+                False,
+                True,
+                False
+            ]),
+            DoesNotRaise()
+        ),  # three boxes with no category
+        (
+            np.array([
+                [0.0, 0.0, 30.0, 40.0, 0.8, 0],
+                [5.0, 5.0, 35.0, 45.0, 0.9, 1],
+                [10.0, 10.0, 40.0, 50.0, 0.85, 2],
+            ]),
+            0.5,
+            np.array([
+                True,
+                True,
+                True
+            ]),
+            DoesNotRaise()
+        ),  # three boxes with same category
+        (
+            np.array([
+                [0.0, 0.0, 30.0, 40.0, 0.8, 0],
+                [5.0, 5.0, 35.0, 45.0, 0.9, 0],
+                [10.0, 10.0, 40.0, 50.0, 0.85, 1],
+            ]),
+            0.5,
+            np.array([
+                False,
+                True,
+                True
+            ]),
+            DoesNotRaise()
+        ),  # three boxes with different category
+    ]
+)
+def test_non_max_suppression(
+        predictions: np.ndarray,
+        iou_threshold: float,
+        expected_result: Optional[np.ndarray],
+        exception: Exception
+) -> None:
+    with exception:
+        result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
+        assert np.array_equal(result, expected_result)
+
+
+@pytest.mark.parametrize(
+    "boxes_xyxy, frame_resolution_wh, expected_result",
+    [
+        (
+            np.empty(shape=(0, 4)),
+            (1280, 720),
+            np.empty(shape=(0, 4)),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [-1.0, 1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [0.0, 1.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, -1.0, 1279.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 0.0, 1279.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1281.0, 719.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1280.0, 719.0]
+            ]),
+        ),
+        (
+            np.array([
+                [1.0, 1.0, 1279.0, 721.0]
+            ]),
+            (1280, 720),
+            np.array([
+                [1.0, 1.0, 1279.0, 720.0]
+            ]),
+        ),
+    ]
+)
+def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
+    result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
+    assert np.array_equal(result, expected_result)