Merge pull request #50 from roboflow/feature/drop_requirement_for_class_id_in_detections

feature/drop_requirement_for_class_id_in_detections

CONTRIBUTING.md

@@ -35,4 +35,17 @@ When creating new functions, please ensure you have the following:
 
 All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary.
 
-PRs must pass all tests and linting requirements before they can be merged.
+PRs must pass all tests and linting requirements before they can be merged.
+
+## 🧹 code quality 
+
+We provide two handy commands inside the `Makefile`, namely:
+
+- `make style` to format the code
+- `make check_code_quality` to check code quality (PEP8 basically)
+
+So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4). 
+
+## 🧪 tests 
+
+[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.

README.md

@@ -2,8 +2,8 @@
   <p>
     <a align="center" href="" target="_blank">
       <img
-        width="850"
-        src="https://media.roboflow.com/open-source/supervision/roboflow-supervision-banner.png?ik-sdk-version=javascript-1.4.3&updatedAt=1674062891088"
+        width="100%"
+        src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
       >
     </a>
   </p>
@@ -62,9 +62,9 @@
 
 ## 👋 hello
 
-A set of easy-to-use utils that will come in handy in any Computer Vision project. **Supervision** is still in 
+A set of easy-to-use utils that will come in handy in any computer vision project. **Supervision is still in 
 pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid 
-and may change.
+and may change.**
 
 ## 💻 install
 
@@ -97,32 +97,30 @@ pip install -e ".[dev]"
 
 Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
 
-## 🎬 videos
+## 🎬 tutorials
 
-Learn how to use YOLOv8, ByteTrack and **Supervision** to detect, track and count objects. 🔥
-[Subscribe](https://www.youtube.com/@Roboflow), and stay up to date with our latest YouTube videos!
+🔥 [Subscribe](https://www.youtube.com/@Roboflow) and stay up to date with the latest changes to the Supervision library. 
 
 <p align="center">
     <a href="https://youtu.be/l_kf9CfZ_8M">
-        <img src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" alt="count-objects-in-zone">
+        <img 
+            width="90%"
+            src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" 
+            alt="count-objects-in-zone"
+        >
     </a>
 </p>
 
 <p align="center">
     <a href="https://youtu.be/OS5qI9YBkfk">
-        <img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="count-objects-crossing-line">
+        <img
+            width="90%"
+            src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" 
+            alt="count-objects-crossing-line"
+        >
     </a>
 </p>
 
-## 🧹 code quality 
+## 🏆 contribution
 
-We provide two handy commands inside the `Makefile`, namely:
-
-- `make style` to format the code
-- `make check_code_quality` to check code quality (PEP8 basically)
-
-So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4). 
-
-## 🧪 tests 
-
-[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
+We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!

docs/detection/annotate.md

@@ -0,0 +1,3 @@
+## BoxAnnotator
+
+:::supervision.detection.annotate.BoxAnnotator

docs/detection/utils.md

@@ -0,0 +1,11 @@
+## 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/detection_utils.md

@@ -1,3 +0,0 @@
-## generate_2d_mask
-
-:::supervision.detection.utils.generate_2d_mask

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

@@ -26,16 +26,17 @@ nav:
   - API reference:
     - Video: video.md
     - Detection:
-        - Core: detection_core.md
-        - Utils: detection_utils.md
-    - Draw: draw.md
-    - Geometry: geometry.md
+        - Core: detection/core.md
+        - Annotate: detection/annotate.md
+        - Utils: detection/utils.md
+    - Draw:
+        - Utils: draw_utils.md
     - Notebook: notebook.md
 
 theme:
   name: 'material'
-  logo: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_white.svg
-  favicon: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_color.svg
+  logo: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671
+  favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671
   palette:
     primary: 'deep purple'
     accent: 'teal'

supervision/__init__.py

@@ -1,6 +1,7 @@
-__version__ = "0.3.0"
+__version__ = "0.3.2"
 
-from supervision.detection.core import BoxAnnotator, Detections
+from supervision.detection.annotate import BoxAnnotator
+from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator
 from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator
 from supervision.detection.utils import generate_2d_mask

supervision/detection/annotate.py

@@ -0,0 +1,116 @@
+from typing import List, Optional, Union
+
+import cv2
+import numpy as np
+
+from supervision.detection.core import Detections
+from supervision.draw.color import Color, ColorPalette
+
+
+class BoxAnnotator:
+    """
+    A class for drawing bounding boxes on an image using detections provided.
+
+    Attributes:
+        color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
+        thickness (int): The thickness of the bounding box lines, default is 2
+        text_color (Color): The color of the text on the bounding box, default is white
+        text_scale (float): The scale of the text on the bounding box, default is 0.5
+        text_thickness (int): The thickness of the text on the bounding box, default is 1
+        text_padding (int): The padding around the text on the bounding box, default is 5
+
+    """
+
+    def __init__(
+        self,
+        color: Union[Color, ColorPalette] = ColorPalette.default(),
+        thickness: int = 2,
+        text_color: Color = Color.black(),
+        text_scale: float = 0.5,
+        text_thickness: int = 1,
+        text_padding: int = 10,
+    ):
+        self.color: Union[Color, ColorPalette] = color
+        self.thickness: int = thickness
+        self.text_color: Color = text_color
+        self.text_scale: float = text_scale
+        self.text_thickness: int = text_thickness
+        self.text_padding: int = text_padding
+
+    def annotate(
+        self,
+        scene: np.ndarray,
+        detections: Detections,
+        labels: Optional[List[str]] = None,
+        skip_label: bool = False,
+    ) -> np.ndarray:
+        """
+        Draws bounding boxes on the frame using the detections provided.
+
+        Parameters:
+            scene (np.ndarray): The image on which the bounding boxes will be drawn
+            detections (Detections): The detections for which the bounding boxes will be drawn
+            labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
+            skip_label (bool): Is set to True, skips bounding box label annotation.
+        Returns:
+            np.ndarray: The image with the bounding boxes drawn on it
+        """
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
+            x1, y1, x2, y2 = xyxy.astype(int)
+            idx = class_id if class_id is not None else i
+            color = (
+                self.color.by_idx(idx)
+                if isinstance(self.color, ColorPalette)
+                else self.color
+            )
+            cv2.rectangle(
+                img=scene,
+                pt1=(x1, y1),
+                pt2=(x2, y2),
+                color=color.as_bgr(),
+                thickness=self.thickness,
+            )
+            if skip_label:
+                continue
+
+            text = (
+                f"{class_id}"
+                if (labels is None or len(detections) != len(labels))
+                else labels[i]
+            )
+
+            text_width, text_height = cv2.getTextSize(
+                text=text,
+                fontFace=font,
+                fontScale=self.text_scale,
+                thickness=self.text_thickness,
+            )[0]
+
+            text_x = x1 + self.text_padding
+            text_y = y1 - self.text_padding
+
+            text_background_x1 = x1
+            text_background_y1 = y1 - 2 * self.text_padding - text_height
+
+            text_background_x2 = x1 + 2 * self.text_padding + text_width
+            text_background_y2 = y1
+
+            cv2.rectangle(
+                img=scene,
+                pt1=(text_background_x1, text_background_y1),
+                pt2=(text_background_x2, text_background_y2),
+                color=color.as_bgr(),
+                thickness=cv2.FILLED,
+            )
+            cv2.putText(
+                img=scene,
+                text=text,
+                org=(text_x, text_y),
+                fontFace=font,
+                fontScale=self.text_scale,
+                color=self.text_color.as_rgb(),
+                thickness=self.text_thickness,
+                lineType=cv2.LINE_AA,
+            )
+        return scene

supervision/detection/core.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Iterator, List, Optional, Tuple, Union
+from typing import Iterator, Optional, Tuple, Union
 
-import cv2
 import numpy as np
 
 from supervision.detection.utils import non_max_suppression
-from supervision.draw.color import Color, ColorPalette
 from supervision.geometry.core import Position
 
 
@@ -18,13 +16,13 @@ class Detections:
 
     Attributes:
         xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
+        class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections.
         confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections.
-        class_id (np.ndarray): An array of shape `(n,)` containing the class ids of the detections.
         tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections.
     """
 
     xyxy: np.ndarray
-    class_id: np.ndarray
+    class_id: Optional[np.ndarray] = None
     confidence: Optional[np.ndarray] = None
     tracker_id: Optional[np.ndarray] = None
 
@@ -32,7 +30,8 @@ class Detections:
         n = len(self.xyxy)
         validators = [
             (isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
-            (isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
+            self.class_id is None
+            or (isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
             self.confidence is None
             or (
                 isinstance(self.confidence, np.ndarray)
@@ -47,8 +46,8 @@ class Detections:
         if not all(validators):
             raise ValueError(
                 "xyxy must be 2d np.ndarray with (n, 4) shape, "
+                "class_id must be None or 1d np.ndarray with (n,) shape, "
                 "confidence must be None or 1d np.ndarray with (n,) shape, "
-                "class_id must be 1d np.ndarray with (n,) shape, "
                 "tracker_id must be None or 1d np.ndarray with (n,) shape"
             )
 
@@ -68,7 +67,7 @@ class Detections:
             yield (
                 self.xyxy[i],
                 self.confidence[i] if self.confidence is not None else None,
-                self.class_id[i],
+                self.class_id[i] if self.class_id is not None else None,
                 self.tracker_id[i] if self.tracker_id is not None else None,
             )
 
@@ -76,13 +75,18 @@ class Detections:
         return all(
             [
                 np.array_equal(self.xyxy, other.xyxy),
+                any(
+                    [
+                        self.class_id is None and other.class_id is None,
+                        np.array_equal(self.class_id, other.class_id),
+                    ]
+                ),
                 any(
                     [
                         self.confidence is None and other.confidence is None,
                         np.array_equal(self.confidence, other.confidence),
                     ]
                 ),
-                np.array_equal(self.class_id, other.class_id),
                 any(
                     [
                         self.tracker_id is None and other.tracker_id is None,
@@ -93,14 +97,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 +117,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 +125,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 +152,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 +166,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 +177,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 +187,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:
@@ -234,8 +217,10 @@ class Detections:
         ):
             return Detections(
                 xyxy=self.xyxy[index],
-                confidence=self.confidence[index],
-                class_id=self.class_id[index],
+                confidence=self.confidence[index]
+                if self.confidence is not None
+                else None,
+                class_id=self.class_id[index] if self.class_id is not None else None,
                 tracker_id=self.tracker_id[index]
                 if self.tracker_id is not None
                 else None,
@@ -246,118 +231,51 @@ 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]
-
-
-class BoxAnnotator:
-    def __init__(
-        self,
-        color: Union[Color, ColorPalette] = ColorPalette.default(),
-        thickness: int = 2,
-        text_color: Color = Color.black(),
-        text_scale: float = 0.5,
-        text_thickness: int = 1,
-        text_padding: int = 10,
-    ):
-        """
-        A class for drawing bounding boxes on an image using detections provided.
-
-        Attributes:
-            color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
-            thickness (int): The thickness of the bounding box lines, default is 2
-            text_color (Color): The color of the text on the bounding box, default is white
-            text_scale (float): The scale of the text on the bounding box, default is 0.5
-            text_thickness (int): The thickness of the text on the bounding box, default is 1
-            text_padding (int): The padding around the text on the bounding box, default is 5
-
-        """
-        self.color: Union[Color, ColorPalette] = color
-        self.thickness: int = thickness
-        self.text_color: Color = text_color
-        self.text_scale: float = text_scale
-        self.text_thickness: int = text_thickness
-        self.text_padding: int = text_padding
-
-    def annotate(
-        self,
-        scene: np.ndarray,
-        detections: Detections,
-        labels: Optional[List[str]] = None,
-        skip_label: bool = False,
-    ) -> np.ndarray:
-        """
-        Draws bounding boxes on the frame using the detections provided.
-
-        Parameters:
-            scene (np.ndarray): The image on which the bounding boxes will be drawn
-            detections (Detections): The detections for which the bounding boxes will be drawn
-            labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
-            skip_label (bool): Is set to True, skips bounding box label annotation.
-        Returns:
-            np.ndarray: The image with the bounding boxes drawn on it
-        """
-        font = cv2.FONT_HERSHEY_SIMPLEX
-        for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
-            x1, y1, x2, y2 = xyxy.astype(int)
-            color = (
-                self.color.by_idx(class_id)
-                if isinstance(self.color, ColorPalette)
-                else self.color
-            )
-            cv2.rectangle(
-                img=scene,
-                pt1=(x1, y1),
-                pt2=(x2, y2),
-                color=color.as_bgr(),
-                thickness=self.thickness,
-            )
-            if skip_label:
-                continue
-
-            text = (
-                f"{class_id}"
-                if (labels is None or len(detections) != len(labels))
-                else labels[i]
-            )
-
-            text_width, text_height = cv2.getTextSize(
-                text=text,
-                fontFace=font,
-                fontScale=self.text_scale,
-                thickness=self.text_thickness,
-            )[0]
-
-            text_x = x1 + self.text_padding
-            text_y = y1 - self.text_padding
-
-            text_background_x1 = x1
-            text_background_y1 = y1 - 2 * self.text_padding - text_height
-
-            text_background_x2 = x1 + 2 * self.text_padding + text_width
-            text_background_y2 = y1
-
-            cv2.rectangle(
-                img=scene,
-                pt1=(text_background_x1, text_background_y1),
-                pt2=(text_background_x2, text_background_y2),
-                color=color.as_bgr(),
-                thickness=cv2.FILLED,
-            )
-            cv2.putText(
-                img=scene,
-                text=text,
-                org=(text_x, text_y),
-                fontFace=font,
-                fontScale=self.text_scale,
-                color=self.text_color.as_rgb(),
-                thickness=self.text_thickness,
-                lineType=cv2.LINE_AA,
-            )
-        return scene

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)