Merge pull request #72 from roboflow/feature/add_pascal_voc_xml_load_support

feature/add_pascal_voc_xml_load_support

docs/annotation/voc.md

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

docs/changelog.md

@@ -1,3 +1,27 @@
+### 0.6.0 <small>April 19, 2023</small>
+
+- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format. 
+- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
+- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
+- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
+- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
+
+### 0.5.2 <small>April 13, 2023</small>
+
+- Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.
+
+### 0.5.1 <small>April 12, 2023</small>
+
+- Fixed `Detections.__getitem__` method did not return mask for selected item.
+- Fixed `Detections.area` crashed for mask detections.
+
+### 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.

docs/dataset/core.md

@@ -0,0 +1,8 @@
+!!! warning
+
+    `Dataset` API is still fluid and may change. If you use Dataset in your project until further notice, freeze the 
+    `supervision` version in your `requirements.txt`.
+
+## Dataset
+
+:::supervision.dataset.core.Dataset

docs/detection/annotate.md

@@ -1,3 +1,7 @@
 ## BoxAnnotator
 
-:::supervision.detection.annotate.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

@@ -1,11 +1,27 @@
-## 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
+:::supervision.detection.utils.non_max_suppression
+
+## polygon_to_mask
+
+:::supervision.detection.utils.polygon_to_mask
+
+## mask_to_xyxy
+
+:::supervision.detection.utils.mask_to_xyxy
+
+## mask_to_polygons
+
+:::supervision.detection.utils.mask_to_polygons
+
+## polygon_to_xyxy
+
+:::supervision.detection.utils.polygon_to_xyxy
+
+## filter_polygons_by_area
+
+:::supervision.detection.utils.filter_polygons_by_area

mkdocs.yml

@@ -24,17 +24,19 @@ extra:
 nav:
   - Home: index.md
   - API reference:
-    - Video: video.md
     - Detection:
         - Core: detection/core.md
         - Annotate: detection/annotate.md
         - Utils: detection/utils.md
+        - Tools:
+            - Polygon Zone: detection/tools/polygon_zone.md
+    - Dataset:
+        - Core: dataset/core.md
     - Draw:
         - Utils: draw/utils.md
-    - Annotations:
-        - Pascal VOC XML: annotation/voc.md
     - Notebook:
         - Utils: notebook/utils.md
+    - Video: video.md
   - Changelog: changelog.md
 
 theme:

supervision/__init__.py

@@ -1,13 +1,22 @@
-__version__ = "0.4.0"
+__version__ = "0.6.0"
 
-from supervision.annotation.voc import detections_to_voc_xml
-from supervision.detection.annotate import BoxAnnotator
+from supervision.dataset.core import Dataset
+from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
 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
+from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
+from supervision.detection.utils import (
+    box_iou_batch,
+    filter_polygons_by_area,
+    mask_to_polygons,
+    mask_to_xyxy,
+    non_max_suppression,
+    polygon_to_mask,
+    polygon_to_xyxy,
+)
 from supervision.draw.color import Color, ColorPalette
 from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
+from supervision.file import list_files_with_extensions
 from supervision.geometry.core import Point, Position, Rect
 from supervision.geometry.utils import get_polygon_center
 from supervision.notebook.utils import plot_image, plot_images_grid

supervision/annotation/voc.py

@@ -1,104 +0,0 @@
-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/dataset/core.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import cv2
+import numpy as np
+
+from supervision.dataset.formats.pascal_voc import (
+    detections_to_pascal_voc,
+    load_pascal_voc_annotations,
+)
+from supervision.detection.core import Detections
+from supervision.file import list_files_with_extensions
+
+
+@dataclass
+class Dataset:
+    """
+    Dataclass containing information about the dataset.
+
+    Attributes:
+        classes (List[str]): List containing dataset class names.
+        images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
+        annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
+    """
+
+    classes: List[str]
+    images: Dict[str, np.ndarray]
+    annotations: Dict[str, Detections]
+
+    def as_pascal_voc(
+        self,
+        images_directory_path: Optional[str] = None,
+        annotations_directory_path: Optional[str] = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.75,
+    ) -> None:
+        """
+        Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
+        PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
+        percentage.
+
+        Args:
+            images_directory_path (Optional[str]): The path to the directory where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path (Optional[str]): The path to the directory where the annotations in
+                PASCAL VOC format should be saved. If not provided, annotations will not be saved.
+            min_image_area_percentage (float): The minimum percentage of detection area relative to
+                the image area for a detection to be included.
+            max_image_area_percentage (float): The maximum percentage of detection area relative to
+                the image area for a detection to be included.
+            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
+        """
+        if images_directory_path:
+            images_path = Path(images_directory_path)
+            images_path.mkdir(parents=True, exist_ok=True)
+
+        if annotations_directory_path:
+            annotations_path = Path(annotations_directory_path)
+            annotations_path.mkdir(parents=True, exist_ok=True)
+
+        for image_name, image in self.images.items():
+            detections = self.annotations[image_name]
+
+            if images_directory_path:
+                cv2.imwrite(str(images_path / image_name), image)
+
+            if annotations_directory_path:
+                annotation_name = Path(image_name).stem
+                pascal_voc_xml = detections_to_pascal_voc(
+                    detections=detections,
+                    classes=self.classes,
+                    filename=image_name,
+                    image_shape=image.shape,
+                    min_image_area_percentage=min_image_area_percentage,
+                    max_image_area_percentage=max_image_area_percentage,
+                    approximation_percentage=approximation_percentage,
+                )
+
+                with open(annotations_path / f"{annotation_name}.xml", "w") as f:
+                    f.write(pascal_voc_xml)
+
+    @classmethod
+    def from_pascal_voc(
+        cls, images_directory_path: str, annotations_directory_path: str
+    ) -> Dataset:
+        """
+        Creates a Dataset instance from PASCAL VOC formatted data.
+
+        Args:
+            images_directory_path (str): The path to the directory containing the images.
+            annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
+
+        Returns:
+            Dataset: A Dataset instance containing the loaded images and annotations.
+        """
+        image_paths = list_files_with_extensions(
+            directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
+        )
+        annotation_paths = list_files_with_extensions(
+            directory=annotations_directory_path, extensions=["xml"]
+        )
+
+        raw_annotations: List[Tuple[str, Detections, List[str]]] = [
+            load_pascal_voc_annotations(annotation_path=str(annotation_path))
+            for annotation_path in annotation_paths
+        ]
+
+        classes = []
+        for annotation in raw_annotations:
+            classes.extend(annotation[2])
+        classes = list(set(classes))
+
+        for annotation in raw_annotations:
+            class_id = [classes.index(class_name) for class_name in annotation[2]]
+            annotation[1].class_id = np.array(class_id)
+
+        images = {
+            image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
+        }
+
+        annotations = {
+            image_name: detections for image_name, detections, _ in raw_annotations
+        }
+        return Dataset(classes=classes, images=images, annotations=annotations)

supervision/dataset/formats/pascal_voc.py

@@ -0,0 +1,171 @@
+from typing import List, Optional, Tuple
+from xml.dom.minidom import parseString
+from xml.etree.ElementTree import Element, SubElement, parse, tostring
+
+import numpy as np
+
+from supervision.detection.core import Detections
+from supervision.detection.utils import (
+    approximate_polygon,
+    filter_polygons_by_area,
+    mask_to_polygons,
+    polygon_to_xyxy,
+)
+
+
+def object_to_pascal_voc(
+    xyxy: np.ndarray, name: str, polygon: Optional[np.ndarray] = None
+) -> Element:
+    root = Element("object")
+
+    object_name = SubElement(root, "name")
+    object_name.text = name
+
+    bndbox = SubElement(root, "bndbox")
+    xmin = SubElement(bndbox, "xmin")
+    xmin.text = str(int(xyxy[0]))
+    ymin = SubElement(bndbox, "ymin")
+    ymin.text = str(int(xyxy[1]))
+    xmax = SubElement(bndbox, "xmax")
+    xmax.text = str(int(xyxy[2]))
+    ymax = SubElement(bndbox, "ymax")
+    ymax.text = str(int(xyxy[3]))
+
+    if polygon is not None:
+        object_polygon = SubElement(root, "polygon")
+        for index, point in enumerate(polygon, start=1):
+            x_coordinate, y_coordinate = point
+            x = SubElement(object_polygon, f"x{index}")
+            x.text = str(x_coordinate)
+            y = SubElement(object_polygon, f"y{index}")
+            y.text = str(y_coordinate)
+
+    return root
+
+
+def detections_to_pascal_voc(
+    detections: Detections,
+    classes: List[str],
+    filename: str,
+    image_shape: Tuple[int, int, int],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> 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.
+        image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections.
+        min_image_area_percentage (float): Minimum detection area relative to area of image associated with it.
+        max_image_area_percentage (float): Maximum detection area relative to area of image associated with it.
+        approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
+    Returns:
+        str: An XML string in Pascal VOC format representing the detections.
+    """
+    height, width, depth = image_shape
+    image_area = height * width
+    minimum_detection_area = min_image_area_percentage * image_area
+    maximum_detection_area = max_image_area_percentage * image_area
+
+    # Create root element
+    annotation = Element("annotation")
+
+    # Add folder element
+    folder = SubElement(annotation, "folder")
+    folder.text = "VOC"
+
+    # Add filename element
+    file_name = SubElement(annotation, "filename")
+    file_name.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 xyxy, mask, _, class_id, _ in detections:
+        name = classes[class_id]
+        if mask is not None:
+            polygons = mask_to_polygons(mask=mask)
+            if len(polygons) == 1:
+                polygons = filter_polygons_by_area(
+                    polygons=polygons, min_area=None, max_area=maximum_detection_area
+                )
+            else:
+                polygons = filter_polygons_by_area(
+                    polygons=polygons,
+                    min_area=minimum_detection_area,
+                    max_area=maximum_detection_area,
+                )
+            for polygon in polygons:
+                approx_polygon = approximate_polygon(
+                    polygon=polygon, percentage=approximation_percentage
+                )
+                xyxy = polygon_to_xyxy(polygon=approx_polygon)
+                next_object = object_to_pascal_voc(
+                    xyxy=xyxy, name=name, polygon=approx_polygon
+                )
+                annotation.append(next_object)
+        else:
+            next_object = object_to_pascal_voc(xyxy=xyxy, name=name)
+            annotation.append(next_object)
+
+    # Generate XML string
+    xml_string = parseString(tostring(annotation)).toprettyxml(indent="  ")
+
+    return xml_string
+
+
+def load_pascal_voc_annotations(
+    annotation_path: str,
+) -> Tuple[str, Detections, List[str]]:
+    """
+    Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names.
+
+    Args:
+        annotation_path (str): The path to the PASCAL VOC XML annotations file.
+
+    Returns:
+        Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections.
+    """
+    tree = parse(annotation_path)
+    root = tree.getroot()
+
+    image_name = root.find("filename").text
+
+    xyxy = []
+    class_names = []
+    for obj in root.findall("object"):
+        class_name = obj.find("name").text
+        class_names.append(class_name)
+
+        bbox = obj.find("bndbox")
+        x1 = int(bbox.find("xmin").text)
+        y1 = int(bbox.find("ymin").text)
+        x2 = int(bbox.find("xmax").text)
+        y2 = int(bbox.find("ymax").text)
+
+        xyxy.append([x1, y1, x2, y2])
+
+    xyxy = np.array(xyxy)
+    detections = Detections(xyxy=xyxy)
+
+    return image_name, detections, class_names

supervision/detection/annotate.py

@@ -56,8 +56,11 @@ class BoxAnnotator:
             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)
+        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)
@@ -114,3 +117,58 @@ class BoxAnnotator:
                 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
+        """
+        if detections.mask is None:
+            return scene
+
+        for i in np.flip(np.argsort(detections.area)):
+            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,55 +1,77 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Iterator, List, Optional, Tuple, Union
+from typing import Any, Iterator, List, Optional, Tuple
 
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression
+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:
     """
     Data class containing information about the detections in a video frame.
-
     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.
+        mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks.
         confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections.
+        class_id (Optional[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: Optional[np.ndarray] = None
+    mask: np.Optional[np.ndarray] = None
     confidence: Optional[np.ndarray] = None
+    class_id: 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)),
-            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)
-                and self.confidence.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, "
-                "class_id must be None or 1d np.ndarray with (n,) shape, "
-                "confidence must be None or 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):
         """
@@ -59,13 +81,22 @@ class Detections:
 
     def __iter__(
         self,
-    ) -> Iterator[Tuple[np.ndarray, Optional[float], int, Optional[Union[str, int]]]]:
+    ) -> 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.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,
@@ -75,6 +106,12 @@ class 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,
@@ -113,7 +150,7 @@ class Detections:
             >>> from supervision import Detections
 
             >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
-            >>> results = model(frame)
+            >>> results = model(IMAGE)
             >>> detections = Detections.from_yolov5(results)
             ```
         """
@@ -141,8 +178,8 @@ class Detections:
             >>> from supervision import Detections
 
             >>> model = YOLO('yolov8s.pt')
-            >>> results = model(frame)[0]
-            >>> detections = Detections.from_yolov8(results)
+            >>> yolov8_results = model(IMAGE)[0]
+            >>> detections = Detections.from_yolov8(yolov8_results)
             ```
         """
         return cls(
@@ -201,6 +238,37 @@ class Detections:
             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 = [], []
@@ -250,6 +318,7 @@ class Detections:
         ):
             return Detections(
                 xyxy=self.xyxy[index],
+                mask=self.mask[index] if self.mask is not None else None,
                 confidence=self.confidence[index]
                 if self.confidence is not None
                 else None,
@@ -264,6 +333,20 @@ class Detections:
 
     @property
     def area(self) -> np.ndarray:
+        """
+        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.
+
+        Returns:
+          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.
+        """
+        if self.mask is not None:
+            return np.array([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.
 

supervision/detection/line_counter.py

@@ -35,7 +35,7 @@ class LineZone:
             detections (Detections): The detections for which to update the counts.
 
         """
-        for xyxy, confidence, class_id, tracker_id in detections:
+        for xyxy, _, confidence, class_id, tracker_id in detections:
             # handle detections with no tracker_id
             if tracker_id is None:
                 continue

supervision/detection/tools/polygon_zone.py

@@ -5,7 +5,7 @@ import cv2
 import numpy as np
 
 from supervision import Detections
-from supervision.detection.utils import clip_boxes, generate_2d_mask
+from supervision.detection.utils import clip_boxes, polygon_to_mask
 from supervision.draw.color import Color
 from supervision.draw.utils import draw_polygon, draw_text
 from supervision.geometry.core import Position
@@ -13,6 +13,17 @@ 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,
@@ -25,11 +36,21 @@ class PolygonZone:
         self.current_count = 0
 
         width, height = frame_resolution_wh
-        self.mask = generate_2d_mask(
+        self.mask = polygon_to_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
         )
@@ -43,6 +64,21 @@ class PolygonZone:
 
 
 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,
@@ -64,6 +100,16 @@ class PolygonZoneAnnotator:
         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,

supervision/detection/utils.py

@@ -1,11 +1,13 @@
-from typing import Tuple
+from typing import List, Optional, Tuple
 
 import cv2
 import numpy as np
 
+MIN_POLYGON_POINT_COUNT = 3
 
-def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
-    """Generate a 2D mask from a polygon.
+
+def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
+    """Generate a mask from a polygon.
 
     Args:
         polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
@@ -115,3 +117,137 @@ def clip_boxes(
     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
+
+
+def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
+    """
+    Converts a binary mask to a list of polygons.
+
+    Parameters:
+        mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`,
+            where H and W are the height and width of the mask, respectively.
+
+    Returns:
+        List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
+            containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
+            are excluded from the output.
+    """
+
+    contours, _ = cv2.findContours(
+        mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+    )
+    return [
+        np.squeeze(contour, axis=1)
+        for contour in contours
+        if contour.shape[0] >= MIN_POLYGON_POINT_COUNT
+    ]
+
+
+def filter_polygons_by_area(
+    polygons: List[np.ndarray],
+    min_area: Optional[float] = None,
+    max_area: Optional[float] = None,
+) -> List[np.ndarray]:
+    """
+    Filters a list of polygons based on their area.
+
+    Parameters:
+        polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
+            containing the `x`, `y` coordinates of the points.
+        min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value
+            will be included in the output. If set to None, no minimum area constraint will be applied.
+        max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value
+            will be included in the output. If set to None, no maximum area constraint will be applied.
+
+    Returns:
+        List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds.
+    """
+    if min_area is None and max_area is None:
+        return polygons
+    ares = [cv2.contourArea(polygon) for polygon in polygons]
+    return [
+        polygon
+        for polygon, area in zip(polygons, ares)
+        if (min_area is None or area >= min_area)
+        and (max_area is None or area <= max_area)
+    ]
+
+
+def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
+    """
+    Converts a polygon represented by a NumPy array into a bounding box.
+
+    Parameters:
+        polygon (np.ndarray): A polygon represented by a NumPy array of shape `(N, 2)`,
+            containing the `x`, `y` coordinates of the points.
+
+    Returns:
+        np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon.
+    """
+    x_min, y_min = np.min(polygon, axis=0)
+    x_max, y_max = np.max(polygon, axis=0)
+    return np.array([x_min, y_min, x_max, y_max])
+
+
+def approximate_polygon(polygon: np.ndarray, percentage: float) -> np.ndarray:
+    """
+    Approximates a given polygon by reducing a certain percentage of points.
+
+    This function uses the Ramer-Douglas-Peucker algorithm to simplify the input polygon by reducing the number of points
+    while preserving the general shape.
+
+    Parameters:
+        polygon (np.ndarray): A 2D NumPy array of shape (N, 2) containing the x, y coordinates of the input polygon's points.
+        percentage (float): The percentage of points to be removed from the input polygon, in the range [0, 1).
+
+    Returns:
+        np.ndarray: A new 2D NumPy array of shape (M, 2), where M <= N * (1 - percentage), containing the x, y coordinates of the
+            approximated polygon's points.
+    """
+
+    if percentage < 0 or percentage >= 1:
+        raise ValueError("Percentage must be in the range [0, 1).")
+
+    target_points = max(int(len(polygon) * (1 - percentage)), 3)
+
+    if len(polygon) <= target_points:
+        return polygon
+
+    epsilon = 0
+    approximated_points = polygon
+    while len(approximated_points) > target_points:
+        epsilon += 0.1
+        approximated_points = cv2.approxPolyDP(polygon, epsilon, closed=True)
+
+    return np.squeeze(approximated_points, axis=1)

supervision/file.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+from typing import List, Optional, Union
+
+
+def write_text_to_file(content: str, output_path: str) -> None:
+    with open(output_path, "w") as file:
+        file.write(content)
+
+
+def list_files_with_extensions(
+    directory: Union[str, Path], extensions: Optional[List[str]] = None
+) -> List[Path]:
+    directory = Path(directory)
+    files_with_extensions = []
+
+    if extensions is not None:
+        for ext in extensions:
+            files_with_extensions.extend(directory.glob(f"*.{ext}"))
+    else:
+        files_with_extensions.extend(directory.glob("*"))
+
+    return files_with_extensions

supervision/notebook/utils.py

@@ -6,7 +6,7 @@ import numpy as np
 
 
 def plot_image(
-    image: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: Optional[str] = "gray"
+    image: np.ndarray, size: Tuple[int, int] = (12, 12), cmap: Optional[str] = "gray"
 ) -> None:
     """
     Plots image using matplotlib.
@@ -27,12 +27,14 @@ def plot_image(
         >>> sv.plot_image(image, (16, 16))
         ```
     """
+    plt.figure(figsize=size)
+
     if image.ndim == 2:
-        plt.figure(figsize=size)
         plt.imshow(image, cmap=cmap)
     else:
-        plt.figure(figsize=size)
         plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+
+    plt.axis("off")
     plt.show()
 
 
@@ -41,6 +43,7 @@ def plot_images_grid(
     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.
@@ -50,6 +53,7 @@ def plot_images_grid(
        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.
@@ -70,7 +74,6 @@ def plot_images_grid(
         >>> plot_images_grid(images, grid_size=(2, 2), titles=titles, figsize=(16, 16))
         ```
     """
-
     nrows, ncols = grid_size
 
     if len(images) > nrows * ncols:
@@ -82,11 +85,13 @@ def plot_images_grid(
 
     for idx, ax in enumerate(axes.flat):
         if idx < len(images):
-            ax.imshow(cv2.cvtColor(images[idx], cv2.COLOR_BGR2RGB))
+            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")
-        else:
-            ax.axis("off")
 
+        ax.axis("off")
     plt.show()

test/detection/test_utils.py

@@ -1,11 +1,11 @@
 from contextlib import ExitStack as DoesNotRaise
-from typing import Optional, Tuple
+from typing import Optional, Tuple, List
 
 import pytest
 
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression, clip_boxes
+from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
 
 
 @pytest.mark.parametrize(
@@ -186,3 +186,93 @@ def test_non_max_suppression(
 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)
+
+
+@pytest.mark.parametrize(
+    "polygons, min_area, max_area, expected_result, exception",
+    [
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            None,
+            None,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # single polygon without area constraints
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            50,
+            None,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # single polygon with min_area constraint
+        (
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            None,
+            50,
+            [],
+            DoesNotRaise()
+        ),  # single polygon with max_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            200,
+            None,
+            [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            None,
+            200,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with max_area constraint
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            200,
+            200,
+            [],
+            DoesNotRaise()
+        ),  # two polygons with both area constraints
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            100,
+            100,
+            [np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area and max_area equal to the area of the first polygon
+        (
+            [
+                np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
+                np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
+            ],
+            400,
+            400,
+            [np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
+            DoesNotRaise()
+        ),  # two polygons with min_area and max_area equal to the area of the second polygon
+    ]
+)
+def test_filter_polygons_by_area(
+        polygons: List[np.ndarray],
+        min_area: Optional[float],
+        max_area: Optional[float],
+        expected_result: List[np.ndarray],
+        exception: Exception
+) -> None:
+    with exception:
+        result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area)
+        assert len(result) == len(expected_result)
+        for result_polygon, expected_result_polygon in zip(result, expected_result):
+            assert np.array_equal(result_polygon, expected_result_polygon)