🛡️ add badges to README.md

feature/dataset_split

README.md

@@ -56,6 +56,10 @@
 
   <br>
 
+[![version](https://badge.fury.io/py/supervision.svg)](https://badge.fury.io/py/supervision)
+[![downloads](https://img.shields.io/pypi/dm/supervision)](https://pypistats.org/packages/supervision)
+[![license](https://img.shields.io/pypi/l/supervision)](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
+[![python-version](https://img.shields.io/pypi/pyversions/supervision)](https://badge.fury.io/py/supervision)
 [![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>

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,30 @@
+### 0.8.0 <small>May 17, 2023</small>
+
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for Dataset inheritance. Current `Dataset` got renamed to `DetectionDataset` and make it inherit from `BaseDataset`.
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using `DetectionDataset.as_yolo`.
+- Added [[#102](https://github.com/roboflow/supervision/pull/103)]: support for splitting `DetectionDataset`.
+- Changed [[#100](https://github.com/roboflow/supervision/pull/100)]: default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
+
+### 0.7.0 <small>May 11, 2023</small>
+
+- Added [[#91](https://github.com/roboflow/supervision/pull/91)]: `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.
+- Added [[#86](https://github.com/roboflow/supervision/pull/86)]: ability to load datasets in YOLO format using `Dataset.from_yolo`. 
+- Added [[#84](https://github.com/roboflow/supervision/pull/84)]: `Detections.merge` to merge multiple `Detections` objects together.
+- Fixed [[#81](https://github.com/roboflow/supervision/pull/81)]: `LineZoneAnnotator.annotate` does not return annotated frame.
+- Changed [[#44](https://github.com/roboflow/supervision/pull/44)]: `LineZoneAnnotator.annotate` to allow for custom text for the in and out tags.
+
+### 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.

docs/dataset/core.md

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

docs/detection/utils.md

@@ -1,7 +1,3 @@
-## generate_2d_mask
-
-:::supervision.detection.utils.generate_2d_mask
-
 ## box_iou_batch
 
 :::supervision.detection.utils.box_iou_batch
@@ -10,6 +6,22 @@
 
 :::supervision.detection.utils.non_max_suppression
 
+## polygon_to_mask
+
+:::supervision.detection.utils.polygon_to_mask
+
 ## mask_to_xyxy
 
-:::supervision.detection.utils.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

docs/quickstart/detections.md

@@ -0,0 +1,312 @@
+## advanced filtering
+
+
+The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down 
+and refine object detections. This section outlines various filtering methods, including filtering by specific class 
+or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones. 
+Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement 
+the filters in their applications.
+
+### by specific class
+
+Allows you to select detections that belong only to one selected class.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.class_id == 0]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-specific-class](https://media.roboflow.com/open-source/supervision/supervision-detection-by-specific-class.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.class_id == 0]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+
+### by set of classes
+
+Allows you to select detections that belong only to selected set of classes. 
+
+=== "After"
+
+    ```python
+    import numpy as np
+    import supervision as sv
+    
+    selected_classes = [0, 2, 3]
+    detections = sv.Detections(...)
+    detections = detections[np.isin(detections.class_id, selected_classes)]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-set-of-classes.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import numpy as np
+    import supervision as sv
+    
+    class_id = [0, 2, 3]
+    detections = sv.Detections(...)
+    detections = detections[np.isin(detections.class_id, class_id)]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by confidence
+
+Allows you to select detections with specific confidence value, for example higher than selected threshold.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.confidence > 0.5]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-confidence.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.confidence > 0.5]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by area
+
+Allows you to select detections based on their size. We define the area as the number of pixels occupied by the 
+detection in the image. In the example below, we have sifted out the detections that are too small. 
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.area > 1000]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-area.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    detections = detections[detections.area > 1000]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by relative area
+
+Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of 
+detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image 
+but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area 
+occupied by them. In the example below, we remove too large detections.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    image = ...
+    height, width, channels = image.shape
+    image_area = height * width
+    
+    detections = sv.Detections(...)
+    detections = detections[(detections.area / image_area) < 0.8]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-relative-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-relative-area.png?updatedAt=1683207183434){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    image = ...
+    height, width, channels = image.shape
+    image_area = height * width
+    
+    detections = sv.Detections(...)
+    detections = detections[(detections.area / image_area) < 0.8]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by box dimensions
+
+Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates, 
+can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively 
+simple and fast.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
+    h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
+    detections = detections[(w > 200) & (h > 200)]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-box-dimensions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-box-dimensions.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    detections = sv.Detections(...)
+    w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
+    h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
+    detections = detections[(w > 200) & (h > 200)]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by `PolygonZone`
+
+Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the 
+zone. In the example below you can see how to filter out all detections located in the lower part of the image.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    zone = sv.PolygonZone(...)
+    detections = sv.Detections(...)
+    mask = zone.trigger(detections=detections)
+    detections = detections[mask]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-polygon-zone](https://media.roboflow.com/open-source/supervision/supervision-detection-by-polygon-zone.png?updatedAt=1683211380445){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    zone = sv.PolygonZone(...)
+    detections = sv.Detections(...)
+    mask = zone.trigger(detections=detections)
+    detections = detections[mask]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>
+
+### by mixed conditions
+
+`Detections`' greatest strength, however, is that you can build arbitrarily complex logical conditions by simply combining separate conditions using `&` or `|`.
+
+=== "After"
+
+    ```python
+    import supervision as sv
+
+    zone = sv.PolygonZone(...)
+    detections = sv.Detections(...)
+    mask = zone.trigger(detections=detections)
+    detections = detections[(detections.confidence > 0.7) & mask]
+    ```
+
+    <div class="result" markdown>
+
+    ![by-mixed-conditions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-mixed-conditions.png){ align=center width="800" }
+
+    </div>
+
+=== "Before"
+
+    ```python
+    import supervision as sv
+
+    zone = sv.PolygonZone(...)
+    detections = sv.Detections(...)
+    mask = zone.trigger(detections=detections)
+    detections = detections[mask]
+    ```
+
+    <div class="result" markdown>
+
+    ![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
+
+    </div>

docs/stylesheets/extra.css

@@ -0,0 +1,4 @@
+:root {
+  --md-primary-fg-color:        #8315F9;
+  --md-accent-fg-color:         #00FFCE;
+}

mkdocs.yml

@@ -21,22 +21,28 @@ extra:
     provider: google
     property: G-P7ZG0Y19G5
 
+extra_css:
+  - stylesheets/extra.css
+
 nav:
   - Home: index.md
+  - Quickstart:
+    - Detections: quickstart/detections.md
   - API reference:
-    - Video: video.md
-    - Detection:
+    - Detections:
         - 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:
+        - Utils: video.md
   - Changelog: changelog.md
 
 theme:
@@ -44,8 +50,19 @@ theme:
   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'
+    # Palette for light mode
+    - scheme: default
+      primary: 'custom'
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+
+    # Palette toggle for dark mode
+    - scheme: slate
+      primary: 'custom'
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
   font:
     text: Roboto
     code: Roboto Mono
@@ -58,5 +75,9 @@ markdown_extensions:
   - admonition
   - pymdownx.details
   - pymdownx.superfences
+  - attr_list
+  - md_in_html
+  - pymdownx.tabbed:
+      alternate_style: true
   - toc:
       permalink: true

setup.py

@@ -26,7 +26,8 @@ setuptools.setup(
     install_requires=[
         'numpy>=1.20.0',
         'opencv-python',
-        'matplotlib'
+        'matplotlib',
+        'pyyaml'
     ],
     packages=find_packages(exclude=("tests",)),
     extras_require={

supervision/__init__.py

@@ -1,13 +1,22 @@
-__version__ = "0.5.1"
+__version__ = "0.8.0"
 
-from supervision.annotation.voc import detections_to_voc_xml
+from supervision.dataset.core import BaseDataset, DetectionDataset
 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.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,332 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterator, 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.dataset.formats.yolo import (
+    load_yolo_annotations,
+    save_data_yaml,
+    save_yolo_annotations,
+)
+from supervision.dataset.ultils import save_dataset_images, train_test_split
+from supervision.detection.core import Detections
+from supervision.file import list_files_with_extensions
+
+
+@dataclass
+class BaseDataset(ABC):
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @abstractmethod
+    def split(
+        self, split_ratio=0.8, random_state=None, shuffle: bool = True
+    ) -> Tuple[BaseDataset, BaseDataset]:
+        pass
+
+
+@dataclass
+class DetectionDataset(BaseDataset):
+    """
+    Dataclass containing information about object detection 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 __len__(self) -> int:
+        """
+        Return the number of images in the dataset.
+
+        Returns:
+            int: The number of images.
+        """
+        return len(self.images)
+
+    def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
+        """
+        Iterate over the images and annotations in the dataset.
+
+        Yields:
+            Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
+                                                          the image data, and its corresponding annotation.
+        """
+        for image_name, image in self.images.items():
+            yield image_name, image, self.annotations.get(image_name, None)
+
+    def split(
+        self, split_ratio=0.8, random_state=None, shuffle: bool = True
+    ) -> Tuple[DetectionDataset, DetectionDataset]:
+        """
+        Splits the dataset into two parts (training and testing) using the provided split_ratio.
+
+        Args:
+            split_ratio (float, optional): The ratio of the training set to the entire dataset. Default is 0.8.
+            random_state (int, optional): The seed for the random number generator. This is used for reproducibility. Default is None.
+            shuffle (bool, optional): Whether to shuffle the data before splitting. Default is True.
+
+        Returns:
+            Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> dataset = sv.DetectionDataset(...)
+            >>> train_dataset, test_dataset = dataset.split(split_ratio=0.7, random_state=42, shuffle=True)
+            >>> len(train_dataset), len(test_dataset)
+            (700, 300)
+            ```
+        """
+
+        image_names = list(self.images.keys())
+        train_names, test_names = train_test_split(
+            data=image_names,
+            train_ratio=split_ratio,
+            random_state=random_state,
+            shuffle=shuffle,
+        )
+
+        train_dataset = DetectionDataset(
+            classes=self.classes,
+            images={name: self.images[name] for name in train_names},
+            annotations={name: self.annotations[name] for name in train_names},
+        )
+        test_dataset = DetectionDataset(
+            classes=self.classes,
+            images={name: self.images[name] for name in test_names},
+            annotations={name: self.annotations[name] for name in test_names},
+        )
+        return train_dataset, test_dataset
+
+    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.0,
+    ) -> 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
+    ) -> DetectionDataset:
+        """
+        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:
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
+
+        Example:
+            ```python
+            >>> import roboflow
+            >>> from roboflow import Roboflow
+            >>> import supervision as sv
+
+            >>> roboflow.login()
+
+            >>> rf = Roboflow()
+
+            >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
+            >>> dataset = project.version(PROJECT_VERSION).download("voc")
+
+            >>> train_dataset = sv.Dataset.from_yolo(
+            ...     images_directory_path=f"{dataset.location}/train/images",
+            ...     annotations_directory_path=f"{dataset.location}/train/labels"
+            ... )
+
+            >>> dataset.classes
+            ['dog', 'person']
+            ```
+        """
+        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 DetectionDataset(classes=classes, images=images, annotations=annotations)
+
+    @classmethod
+    def from_yolo(
+        cls,
+        images_directory_path: str,
+        annotations_directory_path: str,
+        data_yaml_path: str,
+        force_masks: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from YOLO 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 YOLO annotation files.
+            data_yaml_path (str): The path to the data YAML file containing class information.
+            force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
+
+        Returns:
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
+
+        Example:
+            ```python
+            >>> import roboflow
+            >>> from roboflow import Roboflow
+            >>> import supervision as sv
+
+            >>> roboflow.login()
+
+            >>> rf = Roboflow()
+
+            >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
+            >>> dataset = project.version(PROJECT_VERSION).download("yolov5")
+
+            >>> train_dataset = sv.Dataset.from_yolo(
+            ...     images_directory_path=f"{dataset.location}/train/images",
+            ...     annotations_directory_path=f"{dataset.location}/train/labels",
+            ...     data_yaml_path=f"{dataset.location}/data.yaml"
+            ... )
+
+            >>> dataset.classes
+            ['dog', 'person']
+            ```
+        """
+        classes, images, annotations = load_yolo_annotations(
+            images_directory_path=images_directory_path,
+            annotations_directory_path=annotations_directory_path,
+            data_yaml_path=data_yaml_path,
+            force_masks=force_masks,
+        )
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
+
+    def as_yolo(
+        self,
+        images_directory_path: Optional[str] = None,
+        annotations_directory_path: Optional[str] = None,
+        data_yaml_path: Optional[str] = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+    ) -> None:
+        """
+        Exports the dataset to YOLO (You Only Look Once) format. This method saves the images and their corresponding
+        annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
+        for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
+
+        The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
+
+        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
+                YOLO format should be saved. If not provided, annotations will not be saved.
+            data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
+                If not provided, the file 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). This is useful for simplifying the annotations.
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                images_directory_path=images_directory_path, images=self.images
+            )
+        if annotations_directory_path is not None:
+            save_yolo_annotations(
+                annotations_directory_path=annotations_directory_path,
+                images=self.images,
+                annotations=self.annotations,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+        if data_yaml_path is not None:
+            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)

supervision/dataset/formats/pascal_voc.py

@@ -0,0 +1,156 @@
+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.dataset.ultils import approximate_mask_with_polygons
+from supervision.detection.core import Detections
+from supervision.detection.utils import 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
+
+    # 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 = approximate_mask_with_polygons(
+                mask=mask,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+            for polygon in polygons:
+                xyxy = polygon_to_xyxy(polygon=polygon)
+                next_object = object_to_pascal_voc(
+                    xyxy=xyxy, name=name, polygon=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/dataset/formats/yolo.py

@@ -0,0 +1,256 @@
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+import yaml
+
+from supervision.dataset.ultils import approximate_mask_with_polygons
+from supervision.detection.core import Detections
+from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
+from supervision.file import list_files_with_extensions, read_txt_file, save_text_file
+
+
+def _parse_box(values: List[str]) -> np.ndarray:
+    x_center, y_center, width, height = values
+    return np.array(
+        [
+            float(x_center) - float(width) / 2,
+            float(y_center) - float(height) / 2,
+            float(x_center) + float(width) / 2,
+            float(y_center) + float(height) / 2,
+        ],
+        dtype=np.float32,
+    )
+
+
+def _box_to_polygon(box: np.ndarray) -> np.ndarray:
+    return np.array(
+        [[box[0], box[1]], [box[2], box[1]], [box[2], box[3]], [box[0], box[3]]]
+    )
+
+
+def _parse_polygon(values: List[str]) -> np.ndarray:
+    return np.array(values, dtype=np.float32).reshape(-1, 2)
+
+
+def _polygons_to_masks(
+    polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
+) -> np.ndarray:
+    return np.array(
+        [
+            polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
+            for polygon in polygons
+        ],
+        dtype=bool,
+    )
+
+
+def _pair_image_with_annotation(
+    image_paths: List[Union[str, Path]], annotation_paths: List[Union[str, Path]]
+) -> List[Tuple[Union[str, Path], Union[str, Path]]]:
+    image_dict = {p.stem: p for p in image_paths}
+    return [
+        (image_dict[annotation_path.stem], annotation_path)
+        for annotation_path in annotation_paths
+        if annotation_path.stem in image_dict
+    ]
+
+
+def _with_mask(lines: List[str]) -> bool:
+    return any([len(line.split()) > 5 for line in lines])
+
+
+def _extract_class_names(file_path: str) -> List[str]:
+    names = []
+
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+        for line in lines:
+            if line.strip().startswith("names:"):
+                start_index = lines.index(line) + 1
+                break
+
+        for line in lines[start_index:]:
+            if line.strip().startswith("-"):
+                names.append(line.strip().replace("-", "").strip())
+            else:
+                break
+
+    return names
+
+
+def _image_name_to_annotation_name(image_name: str) -> str:
+    base_name, _ = os.path.splitext(image_name)
+    return base_name + ".txt"
+
+
+def yolo_annotations_to_detections(
+    lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
+) -> Detections:
+    if len(lines) == 0:
+        return Detections.empty()
+
+    class_id, relative_xyxy, relative_polygon = [], [], []
+    w, h = resolution_wh
+    for line in lines:
+        values = line.split()
+        class_id.append(int(values[0]))
+        if len(values) == 5:
+            box = _parse_box(values=values[1:])
+            relative_xyxy.append(box)
+            if with_masks:
+                relative_polygon.append(_box_to_polygon(box=box))
+        elif len(values) > 5:
+            polygon = _parse_polygon(values=values[1:])
+            relative_xyxy.append(polygon_to_xyxy(polygon=polygon))
+            if with_masks:
+                relative_polygon.append(polygon)
+
+    class_id = np.array(class_id, dtype=int)
+    relative_xyxy = np.array(relative_xyxy, dtype=np.float32)
+    xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32)
+
+    if not with_masks:
+        return Detections(class_id=class_id, xyxy=xyxy)
+
+    polygons = [
+        (polygon * np.array(resolution_wh)).astype(int) for polygon in relative_polygon
+    ]
+    mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
+    return Detections(class_id=class_id, xyxy=xyxy, mask=mask)
+
+
+def load_yolo_annotations(
+    images_directory_path: str,
+    annotations_directory_path: str,
+    data_yaml_path: str,
+    force_masks: bool = False,
+) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
+    """
+    Loads YOLO annotations and returns class names, images, and their corresponding detections.
+
+    Args:
+        images_directory_path (str): The path to the directory containing the images.
+        annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
+        data_yaml_path (str): The path to the data YAML file containing class information.
+        force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
+
+    Returns:
+        Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values.
+    """
+    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=["txt"]
+    )
+    path_pairs = _pair_image_with_annotation(
+        image_paths=image_paths, annotation_paths=annotation_paths
+    )
+    classes = _extract_class_names(file_path=data_yaml_path)
+    images = {}
+    annotations = {}
+    for image_path, annotation_path in path_pairs:
+        image = cv2.imread(str(image_path))
+        lines = read_txt_file(str(annotation_path))
+        h, w, _ = image.shape
+        resolution_wh = (w, h)
+
+        with_masks = _with_mask(lines=lines)
+        with_masks = force_masks if force_masks else with_masks
+        annotation = yolo_annotations_to_detections(
+            lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
+        )
+
+        images[image_path.name] = image
+        annotations[image_path.name] = annotation
+    return classes, images, annotations
+
+
+def object_to_yolo(
+    xyxy: np.ndarray,
+    class_id: int,
+    image_shape: Tuple[int, int, int],
+    polygon: Optional[np.ndarray] = None,
+) -> str:
+    h, w, _ = image_shape
+    if polygon is None:
+        xyxy_relative = xyxy / np.array([w, h, w, h], dtype=np.float32)
+        x_min, y_min, x_max, y_max = xyxy_relative
+        x_center = (x_min + x_max) / 2
+        y_center = (y_min + y_max) / 2
+        width = x_max - x_min
+        height = y_max - y_min
+        return f"{int(class_id)} {x_center:.5f} {y_center:.5f} {width:.5f} {height:.5f}"
+    else:
+        polygon_relative = polygon / np.array([w, h], dtype=np.float32)
+        polygon_relative = polygon_relative.reshape(-1)
+        polygon_parsed = " ".join([f"{value:.5f}" for value in polygon_relative])
+        return f"{int(class_id)} {polygon_parsed}"
+
+
+def detections_to_yolo_annotations(
+    detections: Detections,
+    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,
+) -> List[str]:
+    annotation = []
+    for xyxy, mask, _, class_id, _ in detections:
+        if mask is not None:
+            polygons = approximate_mask_with_polygons(
+                mask=mask,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+            for polygon in polygons:
+                xyxy = polygon_to_xyxy(polygon=polygon)
+                next_object = object_to_yolo(
+                    xyxy=xyxy,
+                    class_id=class_id,
+                    image_shape=image_shape,
+                    polygon=polygon,
+                )
+                annotation.append(next_object)
+        else:
+            next_object = object_to_yolo(
+                xyxy=xyxy, class_id=class_id, image_shape=image_shape
+            )
+            annotation.append(next_object)
+    return annotation
+
+
+def save_yolo_annotations(
+    annotations_directory_path: str,
+    images: Dict[str, np.ndarray],
+    annotations: Dict[str, Detections],
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> None:
+    Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
+    for image_name, image in images.items():
+        detections = annotations[image_name]
+        yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
+        yolo_annotations_path = os.path.join(
+            annotations_directory_path, yolo_annotations_name
+        )
+        lines = detections_to_yolo_annotations(
+            detections=detections,
+            image_shape=image.shape,
+            min_image_area_percentage=min_image_area_percentage,
+            max_image_area_percentage=max_image_area_percentage,
+            approximation_percentage=approximation_percentage,
+        )
+        save_text_file(lines=lines, file_path=yolo_annotations_path)
+
+
+def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
+    data = {"nc": len(classes), "names": classes}
+    Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
+    with open(data_yaml_path, "w") as outfile:
+        yaml.dump(data, outfile, default_flow_style=False)

supervision/dataset/ultils.py

@@ -0,0 +1,81 @@
+import os
+import random
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, TypeVar
+
+import cv2
+import numpy as np
+
+from supervision.detection.utils import (
+    approximate_polygon,
+    filter_polygons_by_area,
+    mask_to_polygons,
+)
+
+T = TypeVar("T")
+
+
+def approximate_mask_with_polygons(
+    mask: np.ndarray,
+    min_image_area_percentage: float = 0.0,
+    max_image_area_percentage: float = 1.0,
+    approximation_percentage: float = 0.75,
+) -> List[np.ndarray]:
+    height, width = mask.shape
+    image_area = height * width
+    minimum_detection_area = min_image_area_percentage * image_area
+    maximum_detection_area = max_image_area_percentage * image_area
+
+    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,
+        )
+    return [
+        approximate_polygon(polygon=polygon, percentage=approximation_percentage)
+        for polygon in polygons
+    ]
+
+
+def save_dataset_images(
+    images_directory_path: str, images: Dict[str, np.ndarray]
+) -> None:
+    Path(images_directory_path).mkdir(parents=True, exist_ok=True)
+
+    for image_name, image in images.items():
+        target_image_path = os.path.join(images_directory_path, image_name)
+        cv2.imwrite(target_image_path, image)
+
+
+def train_test_split(
+    data: List[T],
+    train_ratio: float = 0.8,
+    random_state: Optional[int] = None,
+    shuffle: bool = True,
+) -> Tuple[List[T], List[T]]:
+    """
+    Splits the data into two parts using the provided train_ratio.
+
+    Args:
+        data (List[T]): The data to split.
+        train_ratio (float): The ratio of the training set to the entire dataset.
+        random_state (Optional[int]): The seed for the random number generator.
+        shuffle (bool): Whether to shuffle the data before splitting.
+
+    Returns:
+        Tuple[List[T], List[T]]: The split data.
+    """
+    if random_state is not None:
+        random.seed(random_state)
+
+    if shuffle:
+        random.shuffle(data)
+
+    split_index = int(len(data) * train_ratio)
+    return data[:split_index], data[split_index:]

supervision/detection/annotate.py

@@ -47,13 +47,34 @@ class BoxAnnotator:
         """
         Draws bounding boxes on the frame using the detections provided.
 
-        Parameters:
+        Args:
             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.
+            labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If `labels` are not provided, corresponding `class_id` will be used as 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
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> classes = ['person', ...]
+            >>> image = ...
+            >>> detections = sv.Detections(...)
+
+            >>> box_annotator = sv.BoxAnnotator()
+            >>> labels = [
+            ...     f"{classes[class_id]} {confidence:0.2f}"
+            ...     for _, _, confidence, class_id, _
+            ...     in detections
+            ... ]
+            >>> annotated_frame = box_annotator.annotate(
+            ...     scene=image.copy(),
+            ...     detections=detections,
+            ...     labels=labels
+            ... )
+            ```
         """
         font = cv2.FONT_HERSHEY_SIMPLEX
         for i in range(len(detections)):
@@ -139,7 +160,7 @@ class MaskAnnotator:
         """
         Overlays the masks on the given image based on the provided detections, with a specified opacity.
 
-        Parameters:
+        Args:
             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
@@ -147,10 +168,10 @@ class MaskAnnotator:
         Returns:
             np.ndarray: The image with the masks overlaid
         """
-        for i in range(len(detections.xyxy)):
-            if detections.mask is None:
-                continue
+        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
             )

supervision/detection/core.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Any, Iterator, List, Optional, Tuple, Union
+from dataclasses import astuple, dataclass
+from typing import Any, Iterator, List, Optional, Tuple
 
 import numpy as np
 
 from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
 from supervision.geometry.core import Position
+from supervision.internal import deprecated
 
 
 def _validate_xyxy(xyxy: Any, n: int) -> None:
@@ -51,19 +52,18 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None:
 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]`
         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.
+        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
     mask: np.Optional[np.ndarray] = None
-    class_id: 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):
@@ -137,7 +137,7 @@ class Detections:
     @classmethod
     def from_yolov5(cls, yolov5_results) -> Detections:
         """
-        Creates a Detections instance from a YOLOv5 output Detections
+        Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result.
 
         Args:
             yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
@@ -147,12 +147,14 @@ class Detections:
 
         Example:
             ```python
+            >>> import cv2
             >>> import torch
-            >>> from supervision import Detections
+            >>> import supervision as sv
 
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
             >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
-            >>> results = model(IMAGE)
-            >>> detections = Detections.from_yolov5(results)
+            >>> result = model(image)
+            >>> detections = sv.Detections.from_yolov5(result)
             ```
         """
         yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
@@ -165,7 +167,7 @@ class Detections:
     @classmethod
     def from_yolov8(cls, yolov8_results) -> Detections:
         """
-        Creates a Detections instance from a YOLOv8 output Results
+        Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
 
         Args:
             yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
@@ -175,12 +177,14 @@ class Detections:
 
         Example:
             ```python
+            >>> import cv2
             >>> from ultralytics import YOLO
-            >>> from supervision import Detections
+            >>> import supervision as sv
 
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
             >>> model = YOLO('yolov8s.pt')
-            >>> yolov8_results = model(IMAGE)[0]
-            >>> detections = Detections.from_yolov8(yolov8_results)
+            >>> result = model(image)[0]
+            >>> detections = sv.Detections.from_yolov8(result)
             ```
         """
         return cls(
@@ -189,10 +193,39 @@ class Detections:
             class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
         )
 
+    @classmethod
+    def from_yolo_nas(cls, yolo_nas_results) -> Detections:
+        """
+        Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result.
+
+        Args:
+            yolo_nas_results (super_gradients.training.models.prediction_results.ImageDetectionPrediction): The output Results instance from YOLO-NAS
+
+        Returns:
+            Detections: A new Detections object.
+
+        Example:
+            ```python
+            >>> import cv2
+            >>> from super_gradients.training import models
+            >>> import supervision as sv
+
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
+            >>> model = models.get('yolo_nas_l', pretrained_weights="coco")
+            >>> result = list(model.predict(image, conf=0.35))[0]
+            >>> detections = sv.Detections.from_yolo_nas(result)
+            ```
+        """
+        return cls(
+            xyxy=yolo_nas_results.prediction.bboxes_xyxy,
+            confidence=yolo_nas_results.prediction.confidence,
+            class_id=yolo_nas_results.prediction.labels.astype(int),
+        )
+
     @classmethod
     def from_transformers(cls, transformers_results: dict) -> Detections:
         """
-        Creates a Detections instance from Object Detection Transformer output Results
+        Creates a Detections instance from object detection [transformer](https://github.com/huggingface/transformers) inference result.
 
         Returns:
             Detections: A new Detections object.
@@ -205,6 +238,32 @@ class Detections:
 
     @classmethod
     def from_detectron2(cls, detectron2_results) -> Detections:
+        """
+        Create a Detections object from the [Detectron2](https://github.com/facebookresearch/detectron2) inference result.
+
+        Args:
+            detectron2_results: The output of a Detectron2 model containing instances with prediction data.
+
+        Returns:
+            (Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
+
+        Example:
+            ```python
+            >>> import cv2
+            >>> from detectron2.engine import DefaultPredictor
+            >>> from detectron2.config import get_cfg
+            >>> import supervision as sv
+
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
+            >>> cfg = get_cfg()
+            >>> cfg.merge_from_file("path/to/config.yaml")
+            >>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
+            >>> predictor = DefaultPredictor(cfg)
+            >>> result = predictor(image)
+
+            >>> detections = sv.Detections.from_detectron2(result)
+            ```
+        """
         return cls(
             xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
             confidence=detectron2_results["instances"].scores.cpu().numpy(),
@@ -216,6 +275,38 @@ class Detections:
 
     @classmethod
     def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
+        """
+        Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result.
+
+        Args:
+            roboflow_result (dict): The result from the Roboflow API containing predictions.
+            class_list (List[str]): A list of class names corresponding to the class IDs in the API result.
+
+        Returns:
+            (Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> roboflow_result = {
+            ...     "predictions": [
+            ...         {
+            ...             "x": 0.5,
+            ...             "y": 0.5,
+            ...             "width": 0.2,
+            ...             "height": 0.3,
+            ...             "class": "person",
+            ...             "confidence": 0.9
+            ...         },
+            ...         # ... more predictions ...
+            ...     ]
+            ... }
+            >>> class_list = ["person", "car", "dog"]
+
+            >>> detections = sv.Detections.from_roboflow(roboflow_result, class_list)
+            ```
+        """
         xyxy = []
         confidence = []
         class_id = []
@@ -242,7 +333,7 @@ class Detections:
     @classmethod
     def from_sam(cls, sam_result: List[dict]) -> Detections:
         """
-        Creates a Detections instance from Segment Anything Model (SAM) by Meta AI.
+        Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result.
 
         Args:
             sam_result (List[dict]): The output Results instance from SAM
@@ -252,8 +343,8 @@ class Detections:
 
         Example:
             ```python
-            >>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
             >>> import supervision as sv
+            >>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 
             >>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
             >>> mask_generator = SamAutomaticMaskGenerator(sam)
@@ -271,6 +362,9 @@ class Detections:
         return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
 
     @classmethod
+    @deprecated(
+        "Dataset loading and saving is going to be executed by supervision.dataset.core.Dataset"
+    )
     def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
         xyxy, class_id = [], []
 
@@ -283,12 +377,74 @@ class Detections:
 
     @classmethod
     def empty(cls) -> Detections:
+        """
+        Create an empty Detections object with no bounding boxes, confidences, or class IDs.
+
+        Returns:
+            (Detections): An empty Detections object.
+
+        Example:
+            ```python
+            >>> from supervision import Detections
+
+            >>> empty_detections = Detections.empty()
+            ```
+        """
         return cls(
             xyxy=np.empty((0, 4), dtype=np.float32),
             confidence=np.array([], dtype=np.float32),
             class_id=np.array([], dtype=int),
         )
 
+    @classmethod
+    def merge(cls, detections_list: List[Detections]) -> Detections:
+        """
+        Merge a list of Detections objects into a single Detections object.
+
+        This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`,
+        `confidence`, `class_id`, and `tracker_id`) into a single Detections object. If all elements in a field are not
+        `None`, the corresponding field will be stacked. Otherwise, the field will be set to `None`.
+
+        Args:
+            detections_list (List[Detections]): A list of Detections objects to merge.
+
+        Returns:
+            (Detections): A single Detections object containing the merged data from the input list.
+
+        Example:
+            ```python
+            >>> from supervision import Detections
+
+            >>> detections_1 = Detections(...)
+            >>> detections_2 = Detections(...)
+
+            >>> merged_detections = Detections.merge([detections_1, detections_2])
+            ```
+        """
+        if len(detections_list) == 0:
+            return Detections.empty()
+
+        detections_tuples_list = [astuple(detection) for detection in detections_list]
+        xyxy, mask, confidence, class_id, tracker_id = [
+            list(field) for field in zip(*detections_tuples_list)
+        ]
+
+        all_not_none = lambda l: all(x is not None for x in l)
+
+        xyxy = np.vstack(xyxy)
+        mask = np.vstack(mask) if all_not_none(mask) else None
+        confidence = np.hstack(confidence) if all_not_none(confidence) else None
+        class_id = np.hstack(class_id) if all_not_none(class_id) else None
+        tracker_id = np.hstack(tracker_id) if all_not_none(tracker_id) else None
+
+        return cls(
+            xyxy=xyxy,
+            mask=mask,
+            confidence=confidence,
+            class_id=class_id,
+            tracker_id=tracker_id,
+        )
+
     def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
         """
         Returns the bounding box coordinates for a specific anchor.
@@ -319,9 +475,7 @@ class Detections:
         ):
             return Detections(
                 xyxy=self.xyxy[index],
-                mask=self.mask[index]
-                if self.mask is not None
-                else None,
+                mask=self.mask[index] if self.mask is not None else None,
                 confidence=self.confidence[index]
                 if self.confidence is not None
                 else None,

supervision/detection/line_counter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional
 
 import cv2
 import numpy as np
@@ -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
@@ -81,6 +81,8 @@ class LineZoneAnnotator:
         text_scale: float = 0.5,
         text_offset: float = 1.5,
         text_padding: int = 10,
+        custom_in_text: Optional[str] = None,
+        custom_out_text: Optional[str] = None,
     ):
         """
         Initialize the LineCounterAnnotator object with default values.
@@ -102,6 +104,8 @@ class LineZoneAnnotator:
         self.text_scale: float = text_scale
         self.text_offset: float = text_offset
         self.text_padding: int = text_padding
+        self.custom_in_text: str = custom_in_text
+        self.custom_out_text: str = custom_out_text
 
     def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray:
         """
@@ -141,8 +145,16 @@ class LineZoneAnnotator:
             lineType=cv2.LINE_AA,
         )
 
-        in_text = f"in: {line_counter.in_count}"
-        out_text = f"out: {line_counter.out_count}"
+        in_text = (
+            f"{self.custom_in_text}: {line_counter.in_count}"
+            if self.custom_in_text is not None
+            else f"in: {line_counter.in_count}"
+        )
+        out_text = (
+            f"{self.custom_out_text}: {line_counter.out_count}"
+            if self.custom_out_text is not None
+            else f"out: {line_counter.out_count}"
+        )
 
         (in_text_width, in_text_height), _ = cv2.getTextSize(
             in_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
@@ -219,3 +231,4 @@ class LineZoneAnnotator:
             self.text_thickness,
             cv2.LINE_AA,
         )
+        return frame

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
@@ -36,7 +36,7 @@ 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)
         )
 

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.
@@ -146,3 +148,113 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
             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, epsilon_step: float = 0.05
+) -> 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)`.
+        epsilon_step (float): Approximation accuracy step. Epsilon is the maximum distance between the original curve and its approximation.
+
+    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 True:
+        epsilon += epsilon_step
+        new_approximated_points = cv2.approxPolyDP(polygon, epsilon, closed=True)
+        if len(new_approximated_points) > target_points:
+            approximated_points = new_approximated_points
+        else:
+            break
+
+    return np.squeeze(approximated_points, axis=1)

supervision/file.py

@@ -0,0 +1,68 @@
+from pathlib import Path
+from typing import List, Optional, Union
+
+
+def list_files_with_extensions(
+    directory: Union[str, Path], extensions: Optional[List[str]] = None
+) -> List[Path]:
+    """
+    List files in a directory with specified extensions or all files if no extensions are provided.
+
+    Args:
+        directory (Union[str, Path]): The directory path as a string or Path object.
+        extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
+
+    Returns:
+        (List[Path]): A list of Path objects for the matching files.
+
+    Examples:
+        ```python
+        >>> from supervision import list_files_with_extensions
+
+        >>> # List all files in the directory
+        >>> files = list_files_with_extensions(directory='my_directory')
+
+        >>> # List only files with '.txt' and '.md' extensions
+        >>> files = list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
+        ```
+    """
+    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
+
+
+def read_txt_file(file_path: str) -> List[str]:
+    """
+    Read a text file and return a list of strings without newline characters.
+
+    Args:
+        file_path (str): The path to the text file.
+
+    Returns:
+        List[str]: A list of strings representing the lines in the text file.
+    """
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+        lines = [line.rstrip("\n") for line in lines]
+
+    return lines
+
+
+def save_text_file(lines: List[str], file_path: str):
+    """
+    Write a list of strings to a text file, each string on a new line.
+
+    Args:
+        lines (List[str]): The list of strings to be written to the file.
+        file_path (str): The path to the text file.
+    """
+    with open(file_path, "w") as file:
+        for line in lines:
+            file.write(line + "\n")

supervision/internal.py

@@ -0,0 +1,18 @@
+import functools
+import warnings
+
+
+def deprecated(reason: str):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            warnings.warn(
+                f"{func.__name__} is deprecated: {reason}",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

supervision/notebook/utils.py

@@ -24,7 +24,7 @@ def plot_image(
         >>> image = cv2.imread("path/to/image.jpg")
 
         %matplotlib inline
-        >>> sv.plot_image(image, (16, 16))
+        >>> sv.plot_image(image=image, size=(16, 16))
         ```
     """
     plt.figure(figsize=size)
@@ -71,7 +71,7 @@ def plot_images_grid(
         >>> titles = ["Image 1", "Image 2", "Image 3"]
 
         %matplotlib inline
-        >>> plot_images_grid(images, grid_size=(2, 2), titles=titles, figsize=(16, 16))
+        >>> plot_images_grid(images, grid_size=(2, 2), titles=titles, size=(16, 16))
         ```
     """
     nrows, ncols = grid_size

test/dataset/formats/test_yolo.py

@@ -0,0 +1,304 @@
+from contextlib import ExitStack as DoesNotRaise
+from typing import List, Tuple, Optional
+
+import pytest
+import numpy as np
+
+from supervision.detection.core import Detections
+from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask, _image_name_to_annotation_name, \
+    object_to_yolo
+
+
+def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
+    x_min, y_min, x_max, y_max = box
+    mask = np.full(resolution_wh, False, dtype=bool)
+    mask[y_min:y_max, x_min:x_max] = True
+    return mask
+
+
+# The result of _mock_simple_mask is a little different from the result produced by cv2.
+def _arrays_almost_equal(arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99) -> bool:
+    equal_elements = np.equal(arr1, arr2)
+    proportion_equal = np.mean(equal_elements)
+    return proportion_equal >= threshold
+
+
+@pytest.mark.parametrize(
+    'lines, expected_result, exception',
+    [
+        (
+            [],
+            False,
+            DoesNotRaise()
+        ),  # empty yolo annotation file
+        (
+            [
+                '0 0.5 0.5 0.2 0.2'
+            ],
+            False,
+            DoesNotRaise()
+        ),  # yolo annotation file with single line with box
+        (
+                [
+                    '0 0.50 0.50 0.20 0.20',
+                    '1 0.11 0.47 0.22 0.30'
+                ],
+                False,
+                DoesNotRaise()
+        ),  # yolo annotation file with two lines with box
+        (
+            [
+                '0 0.5 0.5 0.2 0.2'
+            ],
+            False,
+            DoesNotRaise()
+        ),
+        (
+            [
+                '0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
+            ],
+            True,
+            DoesNotRaise()
+        ),  # yolo annotation file with single line with polygon
+        (
+            [
+                '0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
+                '1 0.11 0.47 0.22 0.30'
+            ],
+            True,
+            DoesNotRaise()
+        ),  # yolo annotation file with two lines - one box and one polygon
+    ]
+)
+def test_with_mask(lines: List[str], expected_result: Optional[bool], exception: Exception) -> None:
+    with exception:
+        result = _with_mask(lines=lines)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'lines, resolution_wh, with_masks, expected_result, exception',
+    [
+        (
+            [],
+            (1000, 1000),
+            False,
+            Detections.empty(),
+            DoesNotRaise()
+        ),  # empty yolo annotation file
+        (
+            [
+                '0 0.5 0.5 0.2 0.2'
+            ],
+            (1000, 1000),
+            False,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600]
+                ], dtype=np.float32),
+                class_id=np.array([0], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with single line with box
+        (
+            [
+                '0 0.50 0.50 0.20 0.20',
+                '1 0.11 0.47 0.22 0.30'
+            ],
+            (1000, 1000),
+            False,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600],
+                    [  0, 320, 220, 620]
+                ], dtype=np.float32),
+                class_id=np.array([0, 1], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with two lines with box
+        (
+            [
+                '0 0.5 0.5 0.2 0.2'
+            ],
+            (1000, 1000),
+            True,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600]
+                ], dtype=np.float32),
+                class_id=np.array([0], dtype=int),
+                mask=np.array([
+                    _mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
+                ], dtype=bool)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with single line with box in with_masks mode
+        (
+            [
+                '0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
+            ],
+            (1000, 1000),
+            True,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600]
+                ], dtype=np.float32),
+                class_id=np.array([0], dtype=int),
+                mask=np.array([
+                    _mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
+                ], dtype=bool)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with single line with polygon
+        (
+            [
+                '0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
+                '1 0.11 0.47 0.22 0.30'
+            ],
+            (1000, 1000),
+            True,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600],
+                    [  0, 320, 220, 620]
+                ], dtype=np.float32),
+                class_id=np.array([0, 1], dtype=int),
+                mask=np.array([
+                    _mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600]),
+                    _mock_simple_mask(resolution_wh=(1000, 1000), box=[  0, 320, 220, 620])
+                ], dtype=bool)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with two lines - one box and one polygon in with_masks mode
+        (
+            [
+                '0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
+                '1 0.11 0.47 0.22 0.30'
+            ],
+            (1000, 1000),
+            False,
+            Detections(
+                xyxy=np.array([
+                    [400, 400, 600, 600],
+                    [  0, 320, 220, 620]
+                ], dtype=np.float32),
+                class_id=np.array([0, 1], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # yolo annotation file with two lines - one box and one polygon
+    ]
+)
+def test_yolo_annotations_to_detections(
+    lines: List[str],
+    resolution_wh: Tuple[int, int],
+    with_masks: bool,
+    expected_result: Optional[Detections],
+    exception: Exception
+) -> None:
+    with exception:
+        result = yolo_annotations_to_detections(
+            lines=lines,
+            resolution_wh=resolution_wh,
+            with_masks=with_masks)
+        assert np.array_equal(result.xyxy, expected_result.xyxy)
+        assert np.array_equal(result.class_id, expected_result.class_id)
+        assert (result.mask is None and expected_result.mask is None) or _arrays_almost_equal(result.mask, expected_result.mask)
+
+
+@pytest.mark.parametrize(
+    'image_name, expected_result, exception',
+    [
+        (
+            'image.png',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple png image
+        (
+            'image.jpeg',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple jpeg image
+        (
+            'image.jpg',
+            'image.txt',
+            DoesNotRaise()
+        ),  # simple jpg image
+        (
+            'image.000.jpg',
+            'image.000.txt',
+            DoesNotRaise()
+        ),  # jpg image with multiple dots in name
+    ]
+)
+def test_image_name_to_annotation_name(
+    image_name: str,
+    expected_result: Optional[str],
+    exception: Exception
+) -> None:
+    with exception:
+        result = _image_name_to_annotation_name(image_name=image_name)
+        assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'xyxy, class_id, image_shape, polygon, expected_result, exception',
+    [
+        (
+            np.array([100, 100, 200, 200], dtype=np.float32),
+            1,
+            (1000, 1000, 3),
+            None,
+            '1 0.15000 0.15000 0.10000 0.10000',
+            DoesNotRaise()
+        ),  # square bounding box on square image
+        (
+            np.array([100, 100, 200, 200], dtype=np.float32),
+            1,
+            (800, 1000, 3),
+            None,
+            '1 0.15000 0.18750 0.10000 0.12500',
+            DoesNotRaise()
+        ),  # square bounding box on horizontal image
+        (
+            np.array([100, 100, 200, 200], dtype=np.float32),
+            1,
+            (1000, 800, 3),
+            None,
+            '1 0.18750 0.15000 0.12500 0.10000',
+            DoesNotRaise()
+        ),  # square bounding box on vertical image
+        (
+            np.array([100, 200, 200, 400], dtype=np.float32),
+            1,
+            (1000, 1000, 3),
+            None,
+            '1 0.15000 0.30000 0.10000 0.20000',
+            DoesNotRaise()
+        ),  # horizontal bounding box on square image
+        (
+            np.array([200, 100, 400, 200], dtype=np.float32),
+            1,
+            (1000, 1000, 3),
+            None,
+            '1 0.30000 0.15000 0.20000 0.10000',
+            DoesNotRaise()
+        ),  # vertical bounding box on square image
+        (
+            np.array([100, 100, 200, 200], dtype=np.float32),
+            1,
+            (1000, 1000, 3),
+            np.array([
+                [100, 100],
+                [200, 100],
+                [200, 200],
+                [100, 100]
+            ], dtype=np.float32),
+            '1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000',
+            DoesNotRaise()
+        ),  # square mask on square image
+    ]
+)
+def test_object_to_yolo(xyxy: np.ndarray, class_id: int, image_shape: Tuple[int, int, int], polygon: Optional[np.ndarray], expected_result: Optional[str], exception: Exception) -> None:
+    with exception:
+        result = object_to_yolo(xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon)
+        assert result == expected_result

test/dataset/test_utils.py

@@ -0,0 +1,82 @@
+from typing import List, TypeVar, Optional, Tuple
+from contextlib import ExitStack as DoesNotRaise
+
+import pytest
+
+from supervision.dataset.ultils import train_test_split
+
+T = TypeVar("T")
+
+
+@pytest.mark.parametrize(
+    'data, train_ratio, random_state, shuffle, expected_result, exception',
+    [
+        (
+            [],
+            0.5,
+            None,
+            False,
+            ([], []),
+            DoesNotRaise()
+        ),  # empty data
+        (
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            0.5,
+            None,
+            False,
+            ([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
+            DoesNotRaise()
+        ),  # data with 10 numbers and 50% train split
+        (
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            1.0,
+            None,
+            False,
+            ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], []),
+            DoesNotRaise()
+        ),  # data with 10 numbers and 100% train split
+        (
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            0.0,
+            None,
+            False,
+            ([], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
+            DoesNotRaise()
+        ),  # data with 10 numbers and 0% train split
+        (
+            ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'],
+            0.5,
+            None,
+            False,
+            (['a', 'b', 'c', 'd', 'e'], ['f', 'g', 'h', 'i', 'j']),
+            DoesNotRaise()
+        ),  # data with 10 chars and 50% train split
+        (
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            0.5,
+            23,
+            True,
+            ([7, 8, 5, 6, 3], [2, 9, 0, 1, 4]),
+            DoesNotRaise()
+        ),  # data with 10 numbers and 50% train split with 23 random seed
+        (
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            0.5,
+            32,
+            True,
+            ([4, 6, 0, 8, 9], [5, 7, 2, 3, 1]),
+            DoesNotRaise()
+        ),  # data with 10 numbers and 50% train split with 23 random seed
+    ]
+)
+def test_train_test_split(
+    data: List[T],
+    train_ratio: float,
+    random_state: int,
+    shuffle: bool,
+    expected_result: Optional[Tuple[List[T], List[T]]],
+    exception: Exception
+) -> None:
+    with exception:
+        result = train_test_split(data=data, train_ratio=train_ratio, random_state=random_state, shuffle=shuffle)
+        assert result == expected_result

test/detection/test_core.py

@@ -4,7 +4,7 @@ import pytest
 
 from supervision import Detections
 
-from typing import Optional, Union
+from typing import Optional, Union, List
 
 import numpy as np
 
@@ -35,6 +35,15 @@ PREDICTIONS = np.array([
     ], dtype=np.float32)
 
 
+def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections:
+    return Detections(
+        xyxy = np.array(xyxy, dtype=np.float32),
+        confidence = confidence if confidence is None else np.array(confidence, dtype=np.float32),
+        class_id = class_id if class_id is None else np.array(class_id, dtype=int),
+        tracker_id = tracker_id if tracker_id is None else np.array(tracker_id, dtype=int),
+    )
+
+
 @pytest.mark.parametrize(
     'detections, index, expected_result, exception',
     [
@@ -103,3 +112,90 @@ def test_getitem(
     with exception:
         result = detections[index]
         assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    'detections_list, expected_result, exception',
+    [
+        (
+            [],
+            Detections.empty(),
+            DoesNotRaise()
+        ),  # empty detections list
+        (
+            [
+                Detections.empty()
+            ],
+            Detections.empty(),
+            DoesNotRaise()
+        ),  # single empty detections
+        (
+            [
+                mock_detections(xyxy=[[10, 10, 20, 20]])
+            ],
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            DoesNotRaise()
+        ),  # single detection with xyxy field
+        (
+            [
+                mock_detections(xyxy=[[10, 10, 20, 20]]),
+                Detections.empty()
+            ],
+            mock_detections(xyxy=[[10, 10, 20, 20]]),
+            DoesNotRaise()
+        ),  # single detection with xyxy field + empty detection
+        (
+            [
+                mock_detections(xyxy=[[10, 10, 20, 20]]),
+                mock_detections(xyxy=[[20, 20, 30, 30]])
+            ],
+            mock_detections(
+                xyxy=[
+                    [10, 10, 20, 20],
+                    [20, 20, 30, 30]
+                ]),
+            DoesNotRaise()
+        ),  # two detections with xyxy field
+        (
+            [
+                mock_detections(
+                    xyxy=[[10, 10, 20, 20]],
+                    class_id=[0]),
+                mock_detections(
+                    xyxy=[[20, 20, 30, 30]])
+            ],
+            mock_detections(
+                xyxy=[
+                    [10, 10, 20, 20],
+                    [20, 20, 30, 30]
+                ]),
+            DoesNotRaise()
+        ),  # detection with xyxy, class_id fields + detection with xyxy field
+(
+            [
+                mock_detections(
+                    xyxy=[[10, 10, 20, 20]],
+                    class_id=[0]),
+                mock_detections(
+                    xyxy=[[20, 20, 30, 30]],
+                    class_id=[1]),
+            ],
+            mock_detections(
+                xyxy=[
+                    [10, 10, 20, 20],
+                    [20, 20, 30, 30]
+                ],
+                class_id=[0, 1]
+            ),
+            DoesNotRaise()
+        ),  # two detections with xyxy, class_id fields
+    ]
+)
+def test_merge(
+        detections_list: List[Detections],
+        expected_result: Optional[Detections],
+        exception: Exception
+) -> None:
+    with exception:
+        result = Detections.merge(detections_list=detections_list)
+        assert result == expected_result

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)