Merge pull request #58 from roboflow/feature/update_to_support_sam

feature/update_to_support_sam

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,11 +2,12 @@
   <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>
+
   <br>
 
   <div align="center">
@@ -53,13 +54,17 @@
       </a>
   </div>
 
+  <br>
+
+[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
+
 </div>
 
 ## 👋 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
 
@@ -92,26 +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/OS5qI9YBkfk">
-        <img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="latest-roboflow-tutorial">
+    <a href="https://youtu.be/l_kf9CfZ_8M">
+        <img 
+            width="90%"
+            src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" 
+            alt="count-objects-in-zone"
+        >
     </a>
 </p>
 
-## 🧹 code quality 
+<p align="center">
+    <a href="https://youtu.be/OS5qI9YBkfk">
+        <img
+            width="90%"
+            src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" 
+            alt="count-objects-crossing-line"
+        >
+    </a>
+</p>
 
-We provide two handy commands inside the `Makefile`, namely:
+## 🏆 contribution
 
-- `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/annotation/voc.md

@@ -0,0 +1,3 @@
+## detections_to_voc_xml
+
+:::supervision.annotation.voc.detections_to_voc_xml

docs/changelog.md

@@ -0,0 +1,40 @@
+### 0.5.0 <small>April 10, 2023</small>
+
+- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.mask` to enable segmentation support.
+- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `MaskAnnotator` to allow easy `Detections.mask` annotation.
+- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
+- Changed [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.area` behaviour to work not only with boxes but also with masks.
+
+### 0.4.0 <small>April 5, 2023</small> 
+
+- Added [[#46](https://github.com/roboflow/supervision/discussions/48)]: `Detections.empty` to allow easy creation of empty `Detections` objects.
+- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
+- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `plot_images_grid` to allow easy plotting of multiple images on single plot.
+- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
+- Changed [[#56](https://github.com/roboflow/supervision/pull/56)]: `show_frame_in_notebook` refactored and renamed to `plot_image`.
+
+### 0.3.2 <small>March 23, 2023</small> 
+
+- Changed [[#50](https://github.com/roboflow/supervision/issues/50)]: Allow `Detections.class_id` to be `None`. 
+
+### 0.3.1 <small>March 6, 2023</small> 
+
+- Fixed [[#41](https://github.com/roboflow/supervision/issues/41)]: `PolygonZone` throws an exception when the object touches the bottom edge of the image.
+- Fixed [[#42](https://github.com/roboflow/supervision/issues/42)]: `Detections.wth_nms` method throws an exception when `Detections` is empty.
+- Changed [[#36](https://github.com/roboflow/supervision/pull/36)]: `Detections.wth_nms` support class agnostic and non-class agnostic case.
+
+### 0.3.0 <small>March 6, 2023</small> 
+
+- Changed: Allow `Detections.confidence` to be `None`.
+- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models. 
+- Added: `Detections.area` to dynamically calculate bounding box area.
+- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation. 
+
+### 0.2.0 <small>February 2, 2023</small> 
+
+- Added: Advanced `Detections` filtering with pandas-like API.
+- Added: `Detections.from_yolov5` and `Detections.from_yolov8` to enable seamless integration with YOLOv5 and YOLOv8 models.
+
+### 0.1.0 <small>January 19, 2023</small> 
+
+Say hello to Supervision 👋

docs/detection/annotate.md

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

docs/detection/tools/polygon_zone.md

@@ -0,0 +1,7 @@
+## PolygonZone
+
+:::supervision.detection.tools.polygon_zone.PolygonZone
+
+## PolygonZoneAnnotator
+
+:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator

docs/detection/utils.md

@@ -0,0 +1,15 @@
+## 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
+
+## mask_to_xyxy
+
+:::supervision.detection.utils.mask_to_xyxy

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

docs/notebook.md

@@ -1,3 +0,0 @@
-## show_frame_in_notebook
-
-:::supervision.notebook.utils.show_frame_in_notebook

docs/notebook/utils.md

@@ -0,0 +1,7 @@
+## plot_image
+
+:::supervision.notebook.utils.plot_image
+
+## plot_images_grid
+
+:::supervision.notebook.utils.plot_images_grid

mkdocs.yml

@@ -26,16 +26,23 @@ nav:
   - API reference:
     - Video: video.md
     - Detection:
-        - Core: detection_core.md
-        - Utils: detection_utils.md
-    - Draw: draw.md
-    - Geometry: geometry.md
-    - Notebook: notebook.md
+        - Core: detection/core.md
+        - Annotate: detection/annotate.md
+        - Utils: detection/utils.md
+        - Tools:
+            - Polygon Zone: detection/tools/polygon_zone.md
+    - Draw:
+        - Utils: draw/utils.md
+    - Annotations:
+        - Pascal VOC XML: annotation/voc.md
+    - Notebook:
+        - Utils: notebook/utils.md
+  - Changelog: changelog.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
+  favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png
   palette:
     primary: 'deep purple'
     accent: 'teal'
@@ -50,4 +57,6 @@ plugins:
 markdown_extensions:
   - admonition
   - pymdownx.details
-  - pymdownx.superfences
+  - pymdownx.superfences
+  - toc:
+      permalink: true

requirements.txt

@@ -1,30 +0,0 @@
-certifi==2022.12.7
-charset-normalizer==3.0.1
-click==8.1.3
-colorama==0.4.6
-ghp-import==2.1.0
-griffe==0.25.4
-idna==3.4
-importlib-metadata==6.0.0
-Jinja2==3.1.2
-Markdown==3.3.7
-MarkupSafe==2.1.2
-mergedeep==1.3.4
-mkdocs==1.4.2
-mkdocs-autorefs==0.4.1
-mkdocs-material==9.0.9
-mkdocs-material-extensions==1.1.1
-mkdocstrings==0.20.0
-mkdocstrings-python==0.8.3
-packaging==23.0
-Pygments==2.14.0
-pymdown-extensions==9.9.2
-python-dateutil==2.8.2
-PyYAML==6.0
-pyyaml_env_tag==0.1
-regex==2022.10.31
-requests==2.28.2
-six==1.16.0
-urllib3==1.26.14
-watchdog==2.2.1
-zipp==3.12.0

setup.py

@@ -24,7 +24,7 @@ setuptools.setup(
     long_description_content_type='text/markdown',
     url='https://github.com/roboflow/supervision',
     install_requires=[
-        'numpy',
+        'numpy>=1.20.0',
         'opencv-python',
         'matplotlib'
     ],

supervision/__init__.py

@@ -1,13 +1,16 @@
-__version__ = "0.2.0"
+__version__ = "0.5.0"
 
-from supervision.detection.core import BoxAnnotator, Detections
-from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator
-from supervision.detection.utils import generate_2d_mask
+from supervision.annotation.voc import detections_to_voc_xml
+from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
+from supervision.detection.core import Detections
+from supervision.detection.line_counter import LineZone, LineZoneAnnotator
+from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
+from supervision.detection.utils import generate_2d_mask, mask_to_xyxy
 from supervision.draw.color import Color, ColorPalette
 from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
 from supervision.geometry.core import Point, Position, Rect
 from supervision.geometry.utils import get_polygon_center
-from supervision.notebook.utils import show_frame_in_notebook
+from supervision.notebook.utils import plot_image, plot_images_grid
 from supervision.video import (
     VideoInfo,
     VideoSink,

supervision/annotation/voc.py

@@ -0,0 +1,104 @@
+from typing import List
+from xml.dom.minidom import parseString
+from xml.etree.ElementTree import Element, SubElement, tostring
+
+from supervision.detection.core import Detections
+
+
+def detections_to_voc_xml(
+    detections: Detections,
+    classes: List[str],
+    filename: str,
+    width: int,
+    height: int,
+    depth: int = 3,
+) -> str:
+    """
+    Converts Detections object to Pascal VOC XML format.
+
+    Args:
+        detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information.
+        classes (List[str]): A list of class names corresponding to the class ids in the Detections object.
+        filename (str): The name of the image file associated with the detections.
+        width (int): The width of the image in pixels.
+        height (int): The height of the image in pixels.
+        depth (int, optional): The number of color channels in the image. Defaults to 3 for RGB images.
+
+    Returns:
+        str: An XML string in Pascal VOC format representing the detections.
+
+    Examples:
+        ```python
+        >>> import numpy as np
+        >>> import supervision as sv
+
+        >>> xyxy = np.array([
+        ...     [50, 30, 200, 180],
+        ...     [20, 40, 150, 190]
+        ... ])
+        >>> class_id = np.array([1, 0])
+        >>> detections = Detections(xyxy=xyxy, class_id=class_id)
+
+        >>> classes = ["dog", "cat"]
+
+        >>> voc_xml = detections_to_voc_xml(
+        ...     detections=detections,
+        ...     classes=classes,
+        ...     filename="image1.jpg",
+        ...     width=500,
+        ...     height=400
+        ... )
+        ```
+    """
+
+    # Create root element
+    annotation = Element("annotation")
+
+    # Add folder element
+    folder = SubElement(annotation, "folder")
+    folder.text = "VOC"
+
+    # Add filename element
+    fname = SubElement(annotation, "filename")
+    fname.text = filename
+
+    # Add source element
+    source = SubElement(annotation, "source")
+    database = SubElement(source, "database")
+    database.text = "roboflow.ai"
+
+    # Add size element
+    size = SubElement(annotation, "size")
+    w = SubElement(size, "width")
+    w.text = str(width)
+    h = SubElement(size, "height")
+    h.text = str(height)
+    d = SubElement(size, "depth")
+    d.text = str(depth)
+
+    # Add segmented element
+    segmented = SubElement(annotation, "segmented")
+    segmented.text = "0"
+
+    # Add object elements
+    for i in range(detections.xyxy.shape[0]):
+        obj = SubElement(annotation, "object")
+
+        class_id = detections.class_id[i] if detections.class_id is not None else None
+        name = SubElement(obj, "name")
+        name.text = classes[class_id] if class_id is not None else "unknown"
+
+        bndbox = SubElement(obj, "bndbox")
+        xmin = SubElement(bndbox, "xmin")
+        xmin.text = str(int(detections.xyxy[i, 0]))
+        ymin = SubElement(bndbox, "ymin")
+        ymin.text = str(int(detections.xyxy[i, 1]))
+        xmax = SubElement(bndbox, "xmax")
+        xmax.text = str(int(detections.xyxy[i, 2]))
+        ymax = SubElement(bndbox, "ymax")
+        ymax.text = str(int(detections.xyxy[i, 3]))
+
+    # Generate XML string
+    xml_string = parseString(tostring(annotation)).toprettyxml(indent="  ")
+
+    return xml_string

supervision/detection/annotate.py

@@ -0,0 +1,174 @@
+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 in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[i].astype(int)
+            class_id = (
+                detections.class_id[i] if detections.class_id is not None else None
+            )
+            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
+
+
+class MaskAnnotator:
+    """
+    A class for overlaying masks on an image using detections provided.
+
+    Attributes:
+        color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette
+    """
+
+    def __init__(
+        self,
+        color: Union[Color, ColorPalette] = ColorPalette.default(),
+    ):
+        self.color: Union[Color, ColorPalette] = color
+
+    def annotate(
+        self, scene: np.ndarray, detections: Detections, opacity: float = 0.5
+    ) -> np.ndarray:
+        """
+        Overlays the masks on the given image based on the provided detections, with a specified opacity.
+
+        Parameters:
+            scene (np.ndarray): The image on which the masks will be overlaid
+            detections (Detections): The detections for which the masks will be overlaid
+            opacity (float): The opacity of the masks, between 0 and 1, default is 0.5
+
+        Returns:
+            np.ndarray: The image with the masks overlaid
+        """
+        for i in range(len(detections.xyxy)):
+            if detections.mask is None:
+                continue
+
+            class_id = (
+                detections.class_id[i] if detections.class_id is not None else None
+            )
+            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
+            )
+
+            mask = detections.mask[i]
+            colored_mask = np.zeros_like(scene, dtype=np.uint8)
+            colored_mask[:] = color.as_bgr()
+
+            scene = np.where(
+                np.expand_dims(mask, axis=-1),
+                np.uint8(opacity * colored_mask + (1 - opacity) * scene),
+                scene,
+            )
+
+        return scene

supervision/detection/core.py

@@ -1,15 +1,52 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Any, Iterator, List, Optional, Tuple, Union
 
-import cv2
 import numpy as np
 
-from supervision.draw.color import Color, ColorPalette
+from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
 from supervision.geometry.core import Position
 
 
+def _validate_xyxy(xyxy: Any, n: int) -> None:
+    is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
+    if not is_valid:
+        raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
+
+
+def _validate_mask(mask: Any, n: int) -> None:
+    is_valid = mask is None or (
+        isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
+    )
+    if not is_valid:
+        raise ValueError("mask must be 3d np.ndarray with (n, W, H) shape")
+
+
+def _validate_class_id(class_id: Any, n: int) -> None:
+    is_valid = class_id is None or (
+        isinstance(class_id, np.ndarray) and class_id.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
+
+
+def _validate_confidence(confidence: Any, n: int) -> None:
+    is_valid = confidence is None or (
+        isinstance(confidence, np.ndarray) and confidence.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
+
+
+def _validate_tracker_id(tracker_id: Any, n: int) -> None:
+    is_valid = tracker_id is None or (
+        isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
+    )
+    if not is_valid:
+        raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
+
+
 @dataclass
 class Detections:
     """
@@ -17,35 +54,25 @@ class Detections:
 
     Attributes:
         xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
-        confidence (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.
+        mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks.
+        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.
         tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections.
     """
 
     xyxy: np.ndarray
-    confidence: np.ndarray
-    class_id: np.ndarray
+    mask: np.Optional[np.ndarray] = None
+    class_id: Optional[np.ndarray] = None
+    confidence: Optional[np.ndarray] = None
     tracker_id: Optional[np.ndarray] = None
 
     def __post_init__(self):
         n = len(self.xyxy)
-        validators = [
-            (isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
-            (isinstance(self.confidence, np.ndarray) and self.confidence.shape == (n,)),
-            (isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
-            self.tracker_id is None
-            or (
-                isinstance(self.tracker_id, np.ndarray)
-                and self.tracker_id.shape == (n,)
-            ),
-        ]
-        if not all(validators):
-            raise ValueError(
-                "xyxy must be 2d np.ndarray with (n, 4) shape, "
-                "confidence must be 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"
-            )
+        _validate_xyxy(xyxy=self.xyxy, n=n)
+        _validate_mask(mask=self.mask, n=n)
+        _validate_class_id(class_id=self.class_id, n=n)
+        _validate_confidence(confidence=self.confidence, n=n)
+        _validate_tracker_id(tracker_id=self.tracker_id, n=n)
 
     def __len__(self):
         """
@@ -53,27 +80,70 @@ class Detections:
         """
         return len(self.xyxy)
 
-    def __iter__(self):
+    def __iter__(
+        self,
+    ) -> Iterator[
+        Tuple[
+            np.ndarray,
+            Optional[np.ndarray],
+            Optional[float],
+            Optional[int],
+            Optional[int],
+        ]
+    ]:
         """
-        Iterates over the Detections object and yield a tuple of `(xyxy, confidence, class_id, tracker_id)` for each detection.
+        Iterates over the Detections object and yield a tuple of `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
         """
         for i in range(len(self.xyxy)):
             yield (
                 self.xyxy[i],
-                self.confidence[i],
-                self.class_id[i],
+                self.mask[i] if self.mask is not None else None,
+                self.confidence[i] if self.confidence is not None else None,
+                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,
             )
 
+    def __eq__(self, other: Detections):
+        return all(
+            [
+                np.array_equal(self.xyxy, other.xyxy),
+                any(
+                    [
+                        self.mask is None and other.mask is None,
+                        np.array_equal(self.mask, other.mask),
+                    ]
+                ),
+                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),
+                    ]
+                ),
+                any(
+                    [
+                        self.tracker_id is None and other.tracker_id is None,
+                        np.array_equal(self.tracker_id, other.tracker_id),
+                    ]
+                ),
+            ]
+        )
+
     @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
@@ -81,11 +151,11 @@ class Detections:
             >>> from supervision import Detections
 
             >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
-            >>> results = model(frame)
+            >>> results = model(IMAGE)
             >>> 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],
@@ -93,14 +163,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
@@ -108,8 +179,8 @@ class Detections:
             >>> from supervision import Detections
 
             >>> model = YOLO('yolov8s.pt')
-            >>> results = model(frame)
-            >>> detections = Detections.from_yolov8(results)
+            >>> yolov8_results = model(IMAGE)[0]
+            >>> detections = Detections.from_yolov8(yolov8_results)
             ```
         """
         return cls(
@@ -118,40 +189,111 @@ class Detections:
             class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
         )
 
-    def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]:
+    @classmethod
+    def from_transformers(cls, transformers_results: dict) -> 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.
+        Creates a Detections instance from Object Detection Transformer output Results
 
         Returns:
-            Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise.
+            Detections: A new Detections object.
         """
-        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,
-            )
+        return cls(
+            xyxy=transformers_results["boxes"].cpu().numpy(),
+            confidence=transformers_results["scores"].cpu().numpy(),
+            class_id=transformers_results["labels"].cpu().numpy().astype(int),
+        )
+
+    @classmethod
+    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(),
+            class_id=detectron2_results["instances"]
+            .pred_classes.cpu()
+            .numpy()
+            .astype(int),
+        )
+
+    @classmethod
+    def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
+        xyxy = []
+        confidence = []
+        class_id = []
+
+        for prediction in roboflow_result["predictions"]:
+            x = prediction["x"]
+            y = prediction["y"]
+            width = prediction["width"]
+            height = prediction["height"]
+            x_min = x - width / 2
+            y_min = y - height / 2
+            x_max = x_min + width
+            y_max = y_min + height
+            xyxy.append([x_min, y_min, x_max, y_max])
+            class_id.append(class_list.index(prediction["class"]))
+            confidence.append(prediction["confidence"])
+
+        return Detections(
+            xyxy=np.array(xyxy),
+            confidence=np.array(confidence),
+            class_id=np.array(class_id).astype(int),
+        )
+
+    @classmethod
+    def from_sam(cls, sam_result: List[dict]) -> Detections:
+        """
+        Creates a Detections instance from Segment Anything Model (SAM) by Meta AI.
+
+        Args:
+            sam_result (List[dict]): The output Results instance from SAM
+
+        Returns:
+            Detections: A new Detections object.
+
+        Example:
+            ```python
+            >>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
+            >>> import supervision as sv
+
+            >>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
+            >>> mask_generator = SamAutomaticMaskGenerator(sam)
+            >>> sam_result = mask_generator.generate(IMAGE)
+            >>> detections = sv.Detections.from_sam(sam_result=sam_result)
+            ```
+        """
+        sorted_generated_masks = sorted(
+            sam_result, key=lambda x: x["area"], reverse=True
+        )
+
+        xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
+        mask = np.array([mask["segmentation"] for mask in sorted_generated_masks])
+
+        return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
+
+    @classmethod
+    def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
+        xyxy, class_id = [], []
+
+        for annotation in coco_annotation:
+            x_min, y_min, width, height = annotation["bbox"]
+            xyxy.append([x_min, y_min, x_min + width, y_min + height])
+            class_id.append(annotation["category_id"])
+
+        return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
+
+    @classmethod
+    def empty(cls) -> Detections:
+        return cls(
+            xyxy=np.empty((0, 4), dtype=np.float32),
+            confidence=np.array([], dtype=np.float32),
+            class_id=np.array([], dtype=int),
+        )
 
     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:
@@ -172,11 +314,15 @@ class Detections:
         raise ValueError(f"{anchor} is not supported.")
 
     def __getitem__(self, index: np.ndarray) -> Detections:
-        if isinstance(index, np.ndarray) and index.dtype == np.bool:
+        if isinstance(index, np.ndarray) and (
+            index.dtype == bool or index.dtype == int
+        ):
             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,
@@ -185,109 +331,67 @@ class Detections:
             f"Detections.__getitem__ not supported for index of type {type(index)}."
         )
 
-
-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,
-    ):
+    @property
+    def area(self) -> np.ndarray:
         """
-        A class for drawing bounding boxes on an image using detections provided.
+        Calculate the area of each detection in the set of object detections. If masks field is defined property
+        returns are of each mask. If only box is given property return area of each box.
 
-        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
+          np.ndarray: An array of floats containing the area of each detection in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
         """
-        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
+        if self.mask is not None:
+            return np.ndarray([np.sum(mask) for mask in self.mask])
+        else:
+            return self.box_area
+
+    @property
+    def box_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, 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."
+
+        if class_agnostic:
+            predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
+            indices = non_max_suppression(
+                predictions=predictions, iou_threshold=threshold
             )
-            cv2.rectangle(
-                img=scene,
-                pt1=(x1, y1),
-                pt2=(x2, y2),
-                color=color.as_bgr(),
-                thickness=self.thickness,
-            )
-            if skip_label:
-                continue
+            return self[indices]
 
-            text = (
-                f"{confidence:0.2f}"
-                if (labels is None or len(detections) != len(labels))
-                else labels[i]
-            )
+        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."
+        )
 
-            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
+        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,80 +0,0 @@
-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.draw.color import Color
-from supervision.draw.utils import draw_polygon, draw_text
-from supervision.geometry.core import Position
-from supervision.geometry.utils import get_polygon_center
-
-
-class PolygonZone:
-    def __init__(
-        self,
-        polygon: np.ndarray,
-        frame_resolution_wh: Tuple[int, int],
-        triggering_position: Position = Position.BOTTOM_CENTER,
-    ):
-        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
-        )
-        is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]]
-        self.current_count = np.sum(is_in_zone)
-        return is_in_zone.astype(bool)
-
-
-class PolygonZoneAnnotator:
-    def __init__(
-        self,
-        zone: PolygonZone,
-        color: Color,
-        thickness: int = 2,
-        text_color: Color = Color.black(),
-        text_scale: float = 0.5,
-        text_thickness: int = 1,
-        text_padding: int = 10,
-    ):
-        self.zone = zone
-        self.color = color
-        self.thickness = thickness
-        self.text_color = text_color
-        self.text_scale = text_scale
-        self.text_thickness = text_thickness
-        self.text_padding = text_padding
-        self.font = cv2.FONT_HERSHEY_SIMPLEX
-        self.center = get_polygon_center(polygon=zone.polygon)
-
-    def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray:
-        annotated_frame = draw_polygon(
-            scene=scene,
-            polygon=self.zone.polygon,
-            color=self.color,
-            thickness=self.thickness,
-        )
-
-        annotated_frame = draw_text(
-            scene=annotated_frame,
-            text=str(self.zone.current_count) if label is None else label,
-            text_anchor=self.center,
-            background_color=self.color,
-            text_color=self.text_color,
-            text_scale=self.text_scale,
-            text_thickness=self.text_thickness,
-            text_padding=self.text_padding,
-            text_font=self.font,
-        )
-
-        return annotated_frame

supervision/detection/tools/polygon_zone.py

@@ -0,0 +1,132 @@
+from dataclasses import replace
+from typing import Optional, Tuple
+
+import cv2
+import numpy as np
+
+from supervision import Detections
+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
+from supervision.geometry.utils import get_polygon_center
+
+
+class PolygonZone:
+    """
+    A class for defining a polygon-shaped zone within a frame for detecting objects.
+
+    Attributes:
+        polygon (np.ndarray): A numpy array defining the polygon vertices
+        frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
+        triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
+        current_count (int): The current count of detected objects within the zone
+        mask (np.ndarray): The 2D bool mask for the polygon zone
+    """
+
+    def __init__(
+        self,
+        polygon: np.ndarray,
+        frame_resolution_wh: Tuple[int, int],
+        triggering_position: Position = Position.BOTTOM_CENTER,
+    ):
+        self.polygon = polygon
+        self.frame_resolution_wh = frame_resolution_wh
+        self.triggering_position = triggering_position
+        self.current_count = 0
+
+        width, height = frame_resolution_wh
+        self.mask = generate_2d_mask(
+            polygon=polygon, resolution_wh=(width + 1, height + 1)
+        )
+
+    def trigger(self, detections: Detections) -> np.ndarray:
+        """
+        Determines if the detections are within the polygon zone.
+
+        Parameters:
+            detections (Detections): The detections to be checked against the polygon zone
+
+        Returns:
+            np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone
+        """
+
+        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)
+
+
+class PolygonZoneAnnotator:
+    """
+    A class for annotating a polygon-shaped zone within a frame with a count of detected objects.
+
+    Attributes:
+        zone (PolygonZone): The polygon zone to be annotated
+        color (Color): The color to draw the polygon lines
+        thickness (int): The thickness of the polygon lines, default is 2
+        text_color (Color): The color of the text on the polygon, default is black
+        text_scale (float): The scale of the text on the polygon, default is 0.5
+        text_thickness (int): The thickness of the text on the polygon, default is 1
+        text_padding (int): The padding around the text on the polygon, default is 10
+        font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX
+        center (Tuple[int, int]): The center of the polygon for text placement
+    """
+
+    def __init__(
+        self,
+        zone: PolygonZone,
+        color: Color,
+        thickness: int = 2,
+        text_color: Color = Color.black(),
+        text_scale: float = 0.5,
+        text_thickness: int = 1,
+        text_padding: int = 10,
+    ):
+        self.zone = zone
+        self.color = color
+        self.thickness = thickness
+        self.text_color = text_color
+        self.text_scale = text_scale
+        self.text_thickness = text_thickness
+        self.text_padding = text_padding
+        self.font = cv2.FONT_HERSHEY_SIMPLEX
+        self.center = get_polygon_center(polygon=zone.polygon)
+
+    def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray:
+        """
+        Annotates the polygon zone within a frame with a count of detected objects.
+
+        Parameters:
+            scene (np.ndarray): The image on which the polygon zone will be annotated
+            label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None)
+
+        Returns:
+            np.ndarray: The image with the polygon zone and count of detected objects
+        """
+        annotated_frame = draw_polygon(
+            scene=scene,
+            polygon=self.zone.polygon,
+            color=self.color,
+            thickness=self.thickness,
+        )
+
+        annotated_frame = draw_text(
+            scene=annotated_frame,
+            text=str(self.zone.current_count) if label is None else label,
+            text_anchor=self.center,
+            background_color=self.color,
+            text_color=self.text_color,
+            text_scale=self.text_scale,
+            text_thickness=self.text_thickness,
+            text_padding=self.text_padding,
+            text_font=self.font,
+        )
+
+        return annotated_frame

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.
 
@@ -18,3 +18,131 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.
     mask = np.zeros((height, width), dtype=np.uint8)
     cv2.fillPoly(mask, [polygon], color=1)
     return mask
+
+
+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.
+
+    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 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
+
+
+def xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray:
+    xyxy = boxes_xywh.copy()
+    xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2]
+    xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3]
+    return xyxy
+
+
+def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
+    """
+    Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.
+
+    Parameters:
+        masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks
+
+    Returns:
+        np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask
+    """
+    n = masks.shape[0]
+    bboxes = np.zeros((n, 4), dtype=int)
+
+    for i, mask in enumerate(masks):
+        rows, cols = np.where(mask)
+
+        if len(rows) > 0 and len(cols) > 0:
+            x_min, x_max = np.min(cols), np.max(cols)
+            y_min, y_max = np.min(rows), np.max(rows)
+            bboxes[i, :] = [x_min, y_min, x_max, y_max]
+
+    return bboxes

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.

supervision/notebook/utils.py

@@ -1,33 +1,97 @@
-from typing import Tuple
+from typing import List, Optional, Tuple
 
 import cv2
 import matplotlib.pyplot as plt
 import numpy as np
 
 
-def show_frame_in_notebook(
-    frame: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: str = "gray"
-):
+def plot_image(
+    image: np.ndarray, size: Tuple[int, int] = (12, 12), cmap: Optional[str] = "gray"
+) -> None:
     """
-    Display a frame in Jupyter Notebook using Matplotlib
+    Plots image using matplotlib.
 
-    Attributes:
-        frame (np.ndarray): The frame to be displayed.
-        size (Tuple[int, int]): The size of the plot. default:(10,10)
-        cmap (str): the colormap to use for single channel images. default:gray
+    Args:
+        image (np.ndarray): The frame to be displayed.
+        size (Tuple[int, int]): The size of the plot.
+        cmap (str): the colormap to use for single channel images.
 
     Examples:
         ```python
-        >>> from supervision.notebook.utils import show_frame_in_notebook
+        >>> import cv2
+        >>> import supervision as sv
+
+        >>> image = cv2.imread("path/to/image.jpg")
 
         %matplotlib inline
-        show_frame_in_notebook(frame, (16, 16))
+        >>> sv.plot_image(image, (16, 16))
         ```
     """
-    if frame.ndim == 2:
-        plt.figure(figsize=size)
-        plt.imshow(frame, cmap=cmap)
+    plt.figure(figsize=size)
+
+    if image.ndim == 2:
+        plt.imshow(image, cmap=cmap)
     else:
-        plt.figure(figsize=size)
-        plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+    plt.axis("off")
+    plt.show()
+
+
+def plot_images_grid(
+    images: List[np.ndarray],
+    grid_size: Tuple[int, int],
+    titles: Optional[List[str]] = None,
+    size: Tuple[int, int] = (12, 12),
+    cmap: Optional[str] = "gray",
+) -> None:
+    """
+    Plots images in a grid using matplotlib.
+
+    Args:
+       images (List[np.ndarray]): A list of images as numpy arrays.
+       grid_size (Tuple[int, int]): A tuple specifying the number of rows and columns for the grid.
+       titles (Optional[List[str]]): A list of titles for each image. Defaults to None.
+       size (Tuple[int, int]): A tuple specifying the width and height of the entire plot in inches.
+       cmap (str): the colormap to use for single channel images.
+
+    Raises:
+       ValueError: If the number of images exceeds the grid size.
+
+    Examples:
+        ```python
+        >>> import cv2
+        >>> import supervision as sv
+
+        >>> image1 = cv2.imread("path/to/image1.jpg")
+        >>> image2 = cv2.imread("path/to/image2.jpg")
+        >>> image3 = cv2.imread("path/to/image3.jpg")
+
+        >>> images = [image1, image2, image3]
+        >>> titles = ["Image 1", "Image 2", "Image 3"]
+
+        %matplotlib inline
+        >>> plot_images_grid(images, grid_size=(2, 2), titles=titles, figsize=(16, 16))
+        ```
+    """
+    nrows, ncols = grid_size
+
+    if len(images) > nrows * ncols:
+        raise ValueError(
+            "The number of images exceeds the grid size. Please increase the grid size or reduce the number of images."
+        )
+
+    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=size)
+
+    for idx, ax in enumerate(axes.flat):
+        if idx < len(images):
+            if images[idx].ndim == 2:
+                ax.imshow(images[idx], cmap=cmap)
+            else:
+                ax.imshow(cv2.cvtColor(images[idx], cv2.COLOR_BGR2RGB))
+
+            if titles is not None and idx < len(titles):
+                ax.set_title(titles[idx])
+
+        ax.axis("off")
     plt.show()

test/detection/test_core.py

@@ -0,0 +1,105 @@
+from contextlib import ExitStack as DoesNotRaise
+
+import pytest
+
+from supervision import Detections
+
+from typing import Optional, Union
+
+import numpy as np
+
+
+PREDICTIONS = np.array([
+        [       2254,         906,        2447,        1353,     0.90538,           0],
+        [       2049,        1133,        2226,        1371,     0.59002,          56],
+        [        727,        1224,         838,        1601,     0.51119,          39],
+        [        808,        1214,         910,        1564,     0.45287,          39],
+        [          6,          52,        1131,        2133,     0.45057,          72],
+        [        299,        1225,         512,        1663,     0.45029,          39],
+        [        529,         874,         645,         945,     0.31101,          39],
+        [          8,          47,        1935,        2135,     0.28192,          72],
+        [       2265,         813,        2328,         901,      0.2714,          62]
+    ], dtype=np.float32)
+
+DETECTIONS = Detections(
+    xyxy=PREDICTIONS[:, :4],
+    confidence=PREDICTIONS[:, 4],
+    class_id=PREDICTIONS[:, 5].astype(int)
+)
+
+
+PREDICTIONS = np.array([
+        [       2254,         906,        2447,        1353,     0.90538,           0],
+        [       2049,        1133,        2226,        1371,     0.59002,          56],
+        [        727,        1224,         838,        1601,     0.51119,          39],
+    ], dtype=np.float32)
+
+
+@pytest.mark.parametrize(
+    'detections, index, expected_result, exception',
+    [
+        (
+            DETECTIONS,
+            DETECTIONS.class_id == 0,
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only detections with class_id = 0
+        (
+            DETECTIONS,
+            DETECTIONS.confidence > 0.5,
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353],
+                    [       2049,        1133,        2226,        1371],
+                    [        727,        1224,         838,        1601]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538,
+                    0.59002,
+                    0.51119
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0,
+                    56,
+                    39
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only detections with confidence > 0.5
+        (
+            DETECTIONS,
+            np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
+            DETECTIONS,
+            DoesNotRaise()
+        ),  # take all detections
+        (
+            DETECTIONS,
+            np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
+            Detections(
+                xyxy=np.empty((0, 4), dtype=np.float32),
+                confidence=np.array([], dtype=np.float32),
+                class_id=np.array([], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take no detections
+    ]
+)
+def test_getitem(
+        detections: Detections,
+        index: Union[int, slice, np.ndarray],
+        expected_result: Optional[Detections],
+        exception: Exception
+) -> None:
+    with exception:
+        result = detections[index]
+        assert result == expected_result

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)