Merge pull request #136 from roboflow/0.10.0-release

`0.10.0` release

.gitattributes

@@ -0,0 +1 @@
+*.ipynb linguist-vendored

CITATION.cff

@@ -2,7 +2,7 @@ cff-version: 1.2.0
 author: Roboflow
 message: If you use this software, please cite it as below.
 title: supervision
-version: 0.1.0
+version: 0.10.0
 date-released: 2023-01-19
 license: MIT
 repository-code: https://github.com/roboflow/supervision

README.md

@@ -10,6 +10,233 @@
 
   <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/supervision/blob/main/demo.ipynb)
+
+</div>
+
+## 👋 hello
+
+**We write your reusable computer vision tools.** Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us! 🤝
+
+## 💻 install
+
+Pip install the supervision package in a
+[**3.11>=Python>=3.7**](https://www.python.org/) environment.
+
+```bash
+pip install supervision
+```
+
+<details close>
+<summary>👉 install from source</summary>
+
+```bash
+# clone repository and navigate to root directory
+git clone https://github.com/roboflow/supervision.git
+cd supervision
+
+# setup python environment and activate it
+python3 -m venv venv
+source venv/bin/activate
+
+# install
+pip install -e ".[dev]"
+```
+
+</details>
+
+## 🔥 quickstart
+
+### [detections processing](https://roboflow.github.io/supervision/detection/core/)
+
+```python
+>>> import supervision as sv
+>>> from ultralytics import YOLO
+
+>>> model = YOLO('yolov8s.pt')
+>>> result = model(IMAGE)[0]
+>>> detections = sv.Detections.from_yolov8(result)
+
+>>> len(detections)
+5
+```
+
+<details close>
+<summary>👉 more detections utils</summary>
+  
+- Easily switch inference pipeline between supported object detection / instance segmentation models
+  
+    ```python
+    >>> 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)
+    >>> sam_result = mask_generator.generate(IMAGE)
+    >>> detections = sv.Detections.from_sam(sam_result=sam_result)
+    ```
+ 
+- [Advanced filtering](https://roboflow.github.io/supervision/quickstart/detections/)
+  
+    ```python
+    >>> detections = detections[detections.class_id == 0]
+    >>> detections = detections[detections.confidence > 0.5]
+    >>> detections = detections[detections.area > 1000]
+    ```
+  
+- Image annotation
+  
+    ```python
+    >>> import supervision as sv
+
+    >>> box_annotator = sv.BoxAnnotator()
+    >>> annotated_frame = box_annotator.annotate(
+    ...     scene=IMAGE,
+    ...     detections=detections
+    ... )
+    ```
+  
+</details>
+
+### [datasets processing](https://roboflow.github.io/supervision/dataset/core/)
+
+```python
+>>> import supervision as sv
+
+>>> dataset = sv.DetectionDataset.from_yolo(
+...     images_directory_path='...',
+...     annotations_directory_path='...',
+...     data_yaml_path='...'
+... )
+
+>>> dataset.classes
+['dog', 'person']
+
+>>> len(dataset)
+1000
+```
+
+<details close>
+<summary>👉 more dataset utils</summary>
+
+- Load object detection / instance segmentation datasets in one of supported formats
+
+    ```python
+    >>> dataset = sv.DetectionDataset.from_yolo(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...',
+    ...     data_yaml_path='...'
+    ... )
+
+    >>> dataset = sv.DetectionDataset.from_pascal_voc(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...'
+    ... )
+    ```
+  
+- Loop over dataset entries
+
+    ```python
+    >>> for name, image, labels in dataset:
+    ...     print(labels.xyxy)
+
+    array([[404.      , 719.      , 538.      , 884.5     ],
+           [155.      , 497.      , 404.      , 833.5     ],
+           [ 20.154999, 347.825   , 416.125   , 915.895   ]], dtype=float32)
+    ```
+  
+- Split dataset for training, testing and validation
+  
+    ```python
+    >>> train_dataset, test_dataset = dataset.split(split_ratio=0.7)
+    >>> test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
+  
+    >>> len(train_dataset), len(test_dataset), len(valid_dataset)
+    (700, 150, 150)
+    ```
+  
+- Save object detection / instance segmentation datasets in one of supported formats
+  
+    ```python
+    >>> dataset.as_yolo(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...',
+    ...     data_yaml_path='...'
+    ... )
+
+    >>> dataset.as_pascal_voc(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...'
+    ... )
+    ```
+  
+- Convert labels between suppoted formats
+  
+    ```python
+    >>> sv.DetectionDataset.from_yolo(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...',
+    ...     data_yaml_path='...'
+    ... ).as_pascal_voc(
+    ...     images_directory_path='...',
+    ...     annotations_directory_path='...'
+    ... )
+    ```
+  
+- Load classification datasets in one of supported formats
+
+    ```python
+    >>> cs = sv.ClassificationDataset.from_folder_structure(
+    ...     root_directory_path='...'
+    ... )
+    ```
+
+- Save classification datasets in one of supported formats
+
+    ```python
+    >>> cs.as_folder_structure(
+    ...     root_directory_path='...'
+    ... )
+    ```
+
+</details>
+
+## 🎬 tutorials
+
+<p align="left">
+<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
+<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
+<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
+<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p> 
+
+<br/> 
+
+<p align="left">
+<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
+<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
+<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
+<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
+
+## 📚 documentation
+
+Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
+
+## 💜 built with supervision
+
+You built something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
+
+## 🏆 contribution
+
+We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!
+
+<br>
+
+<div align="center">
+
   <div align="center">
       <a href="https://youtube.com/roboflow">
           <img
@@ -54,73 +281,4 @@
       </a>
   </div>
 
-  <br>
-
-[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
-
 </div>
-
-## 👋 hello
-
-A set of easy-to-use utils that will come in handy in any computer vision project. **Supervision is still in 
-pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid 
-and may change.**
-
-## 💻 install
-
-Pip install the supervision package in a
-[**3.10>=Python>=3.7**](https://www.python.org/) environment.
-
-```bash
-pip install supervision
-```
-
-<details close>
-<summary>Install from source</summary>
-
-```bash
-# clone repository and navigate to root directory
-git clone https://github.com/roboflow/supervision.git
-cd supervision
-
-# setup python environment and activate it
-python3 -m venv venv
-source venv/bin/activate
-
-# install
-pip install -e ".[dev]"
-```
-
-</details>
-
-## 📖 documentation
-
-Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
-
-## 🎬 tutorials
-
-🔥 [Subscribe](https://www.youtube.com/@Roboflow) and stay up to date with the latest changes to the Supervision library. 
-
-<p align="center">
-    <a href="https://youtu.be/l_kf9CfZ_8M">
-        <img 
-            width="90%"
-            src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" 
-            alt="count-objects-in-zone"
-        >
-    </a>
-</p>
-
-<p align="center">
-    <a href="https://youtu.be/OS5qI9YBkfk">
-        <img
-            width="90%"
-            src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" 
-            alt="count-objects-crossing-line"
-        >
-    </a>
-</p>
-
-## 🏆 contribution
-
-We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!

docs/changelog.md

@@ -1,9 +1,107 @@
+### 0.10.0 <small>June 14, 2023</small>
+
+- Added [[#125](https://github.com/roboflow/supervision/pull/125)]: ability to load and save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) in a folder structure format.
+
+```python
+>>> import supervision as sv
+
+>>> cs = sv.ClassificationDataset.from_folder_structure(
+...     root_directory_path='...'
+... )
+
+>>> cs.as_folder_structure(
+...     root_directory_path='...'
+... )
+```
+
+- Added [[#125](https://github.com/roboflow/supervision/pull/125)]: support for [`sv.ClassificationDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.split) allowing to divide `sv.ClassificationDataset` into two parts.
+
+- Added [[#110](https://github.com/roboflow/supervision/pull/110)]: ability to extract masks from Roboflow API results using [`sv.Detections.from_roboflow`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_roboflow). 
+
+- Added [[commit hash](https://github.com/roboflow/supervision/commit/d000292eb2f2342544e0947b65528082e60fb8d6)]: Supervision Quickstart [notebook](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb) where you can learn more about Detection, Dataset and Video APIs.
+
+- Changed [[#135](https://github.com/roboflow/supervision/pull/135)]: `sv.get_video_frames_generator` documentation to better describe actual behavior.  
+
+### 0.9.0 <small>June 7, 2023</small>
+
+- Added [[#118](https://github.com/roboflow/supervision/pull/118)]: ability to select [`sv.Detections`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.__getitem__) by index, list of indexes or slice. Here is an example illustrating the new selection methods.
+
+```python
+>>> import supervision as sv
+
+>>> detections = sv.Detections(...)
+>>> len(detections[0])
+1
+>>> len(detections[[0, 1]])
+2
+>>> len(detections[0:2])
+2
+```
+
+- Added [[#101](https://github.com/roboflow/supervision/pull/101)]: ability to extract masks from YOLOv8 result using [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8). Here is an example illustrating how to extract boolean masks from the result of the YOLOv8 model inference.
+ 
+- Added [[#122](https://github.com/roboflow/supervision/pull/122)]: ability to crop image using [`sv.crop`](https://roboflow.github.io/supervision/utils/image/#crop). Here is an example showing how to get a separate crop for each detection in `sv.Detections`.
+
+- Added [[#120](https://github.com/roboflow/supervision/pull/120)]: ability to conveniently save multiple images into directory using [`sv.ImageSink`](https://roboflow.github.io/supervision/utils/image/#imagesink). Here is an example showing how to save every tenth video frame as a separate image.
+
+```python
+>>> import supervision as sv
+
+>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
+...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=10):
+...         sink.save_image(image=image)
+```
+
+- Fixed [[#106](https://github.com/roboflow/supervision/issues/106)]: inconvenient handling of [`sv.PolygonZone`](https://roboflow.github.io/supervision/detection/tools/polygon_zone/#polygonzone) coordinates. Now `sv.PolygonZone` accepts coordinates in the form of `[[x1, y1], [x2, y2], ...]` that can be both integers and floats.
+
+### 0.8.0 <small>May 17, 2023</small>
+
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for dataset inheritance. The current `Dataset` got renamed to `DetectionDataset`. Now [`DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) inherits from `BaseDataset`. This change was made to enforce the future consistency of APIs of different types of computer vision datasets.
+- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo). 
+
+```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")
+
+>>> ds = sv.DetectionDataset.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"
+... )
+
+>>> ds.classes
+['dog', 'person']
+```
+
+- Added [[#102](https://github.com/roboflow/supervision/pull/103)]: support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts. 
+
+```python
+>>> import supervision as sv
+
+>>> ds = sv.DetectionDataset(...)
+>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
+
+>>> len(train_ds), len(test_ds)
+(700, 300)
+```
+
+- 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>
 

docs/classification/core.md

@@ -0,0 +1,3 @@
+## Classifications
+
+:::supervision.classification.core.Classifications

docs/dataset/core.md

@@ -1,8 +1,12 @@
 !!! warning
 
-    `Dataset` API is still fluid and may change. If you use Dataset in your project until further notice, freeze the 
-    `supervision` version in your `requirements.txt`.
+    Dataset 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`.
 
-## Dataset
+## DetectionDataset
 
-:::supervision.dataset.core.Dataset
+:::supervision.dataset.core.DetectionDataset
+
+## ClassificationDataset
+
+:::supervision.dataset.core.ClassificationDataset

docs/index.md

@@ -19,7 +19,7 @@ pre-release stage 🚧 Keep your eyes open for potential bugs and be aware that
 ## 💻 How to Install
 
 You can install `supervision` with pip in a 
-[**3.10>=Python>=3.7**](https://www.python.org/) environment.
+[**3.11>=Python>=3.7**](https://www.python.org/) environment.
 
 !!! example "Pip install method (recommended)"
 

docs/notebook/utils.md

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

docs/utils/file.md

@@ -0,0 +1,3 @@
+## list_files_with_extensions
+
+:::supervision.utils.file.list_files_with_extensions

docs/utils/image.md

@@ -0,0 +1,7 @@
+## ImageSink
+
+:::supervision.utils.image.ImageSink
+
+## crop
+
+:::supervision.utils.image.crop

docs/utils/notebook.md

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

docs/utils/video.md

@@ -0,0 +1,15 @@
+## VideoInfo
+
+:::supervision.utils.video.VideoInfo
+
+## VideoSink
+
+:::supervision.utils.video.VideoSink
+
+## get_video_frames_generator
+
+:::supervision.utils.video.get_video_frames_generator
+
+## process_video
+
+:::supervision.utils.video.process_video

docs/video.md

@@ -1,15 +0,0 @@
-## VideoInfo
-
-:::supervision.video.VideoInfo
-
-## VideoSink
-
-:::supervision.video.VideoSink
-
-## get_video_frames_generator
-
-:::supervision.video.get_video_frames_generator
-
-## process_video
-
-:::supervision.video.process_video

mkdocs.yml

@@ -29,6 +29,8 @@ nav:
   - Quickstart:
     - Detections: quickstart/detections.md
   - API reference:
+    - Classifications:
+        - Core: classification/core.md
     - Detections:
         - Core: detection/core.md
         - Annotate: detection/annotate.md
@@ -39,10 +41,11 @@ nav:
         - Core: dataset/core.md
     - Draw:
         - Utils: draw/utils.md
-    - Notebook:
-        - Utils: notebook/utils.md
-    - Video:
-        - Utils: video.md
+    - Utils:
+        - Video: utils/video.md
+        - Image: utils/image.md
+        - Notebook: utils/notebook.md
+        - File: utils/file.md
   - Changelog: changelog.md
 
 theme:

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={
@@ -51,6 +52,7 @@ setuptools.setup(
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',
         'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
         'Programming Language :: Python :: 3 :: Only',
         'Topic :: Software Development',
         'Topic :: Scientific/Engineering',
@@ -62,6 +64,6 @@ setuptools.setup(
         'Operating System :: Unix',
         'Operating System :: MacOS'
     ],
-    keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, Roboflow",
+    keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, SAM, Roboflow",
     python_requires='>=3.7',
 )

supervision/__init__.py

@@ -1,6 +1,11 @@
-__version__ = "0.7.0"
+__version__ = "0.10.0"
 
-from supervision.dataset.core import Dataset
+from supervision.classification.core import Classifications
+from supervision.dataset.core import (
+    BaseDataset,
+    ClassificationDataset,
+    DetectionDataset,
+)
 from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
 from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator
@@ -16,11 +21,12 @@ from supervision.detection.utils import (
 )
 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
-from supervision.video import (
+from supervision.utils.file import list_files_with_extensions
+from supervision.utils.image import ImageSink, crop
+from supervision.utils.notebook import plot_image, plot_images_grid
+from supervision.utils.video import (
     VideoInfo,
     VideoSink,
     get_video_frames_generator,

supervision/classification/core.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+
+def _validate_class_ids(class_id: Any, n: int) -> None:
+    """
+    Ensure that class_id is a 1d np.ndarray with (n, ) shape.
+    """
+    is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
+    if not is_valid:
+        raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")
+
+
+def _validate_confidence(confidence: Any, n: int) -> None:
+    """
+    Ensure that confidence is a 1d np.ndarray with (n, ) shape.
+    """
+    if confidence is not None:
+        is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
+        if not is_valid:
+            raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")
+
+
+@dataclass
+class Classifications:
+    class_id: np.ndarray
+    confidence: Optional[np.ndarray] = None
+
+    def __post_init__(self) -> None:
+        """
+        Validate the classification inputs.
+        """
+        n = len(self.class_id)
+
+        _validate_class_ids(self.class_id, n)
+        _validate_confidence(self.confidence, n)
+
+    @classmethod
+    def from_yolov8(cls, yolov8_results) -> Classifications:
+        """
+        Creates a Classifications 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
+
+        Returns:
+            Detections: A new Classifications object.
+
+        Example:
+            ```python
+            >>> import cv2
+            >>> from ultralytics import YOLO
+            >>> import supervision as sv
+
+            >>> image = cv2.imread(SOURCE_IMAGE_PATH)
+            >>> model = YOLO('yolov8s-cls.pt')
+            >>> result = model(image)[0]
+            >>> classifications = sv.Classifications.from_yolov8(result)
+            ```
+        """
+        confidence = yolov8_results.probs.data.cpu().numpy()
+        return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
+
+    def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Retrieve the top k class IDs and confidences, ordered in descending order by confidence.
+
+        Args:
+            k (int): The number of top class IDs and confidences to retrieve.
+
+        Returns:
+            Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> classifications = sv.Classifications(...)
+
+            >>> classifications.get_top_k(1)
+
+            (array([1]), array([0.9]))
+            ```
+        """
+        if self.confidence is None:
+            raise ValueError("top_k could not be calculated, confidence is None")
+
+        order = np.argsort(self.confidence)[::-1]
+        top_k_order = order[:k]
+        top_k_class_id = self.class_id[top_k_order]
+        top_k_confidence = self.confidence[top_k_order]
+
+        return top_k_class_id, top_k_confidence

supervision/dataset/core.py

@@ -1,25 +1,46 @@
 from __future__ import annotations
 
+import os
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Iterator
+from typing import Dict, Iterator, List, Optional, Tuple
 
 import cv2
 import numpy as np
 
+from supervision.classification.core import Classifications
 from supervision.dataset.formats.pascal_voc import (
     detections_to_pascal_voc,
     load_pascal_voc_annotations,
 )
-from supervision.dataset.formats.yolo import load_yolo_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
+from supervision.utils.file import list_files_with_extensions
 
 
 @dataclass
-class Dataset:
+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 the dataset.
+    Dataclass containing information about object detection dataset.
 
     Attributes:
         classes (List[str]): List containing dataset class names.
@@ -51,13 +72,58 @@ class Dataset:
         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.
+            random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
+            shuffle (bool, optional): Whether to shuffle the data before splitting.
+
+        Returns:
+            Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> ds = sv.DetectionDataset(...)
+            >>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
+            >>> len(train_ds), len(test_ds)
+            (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.75,
+        approximation_percentage: float = 0.0,
     ) -> None:
         """
         Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
@@ -107,7 +173,7 @@ class Dataset:
     @classmethod
     def from_pascal_voc(
         cls, images_directory_path: str, annotations_directory_path: str
-    ) -> Dataset:
+    ) -> DetectionDataset:
         """
         Creates a Dataset instance from PASCAL VOC formatted data.
 
@@ -116,7 +182,7 @@ class Dataset:
             annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
 
         Returns:
-            Dataset: A Dataset instance containing the loaded images and annotations.
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
 
         Example:
             ```python
@@ -131,12 +197,12 @@ class Dataset:
             >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
             >>> dataset = project.version(PROJECT_VERSION).download("voc")
 
-            >>> train_dataset = sv.Dataset.from_yolo(
+            >>> ds = sv.DetectionDataset.from_yolo(
             ...     images_directory_path=f"{dataset.location}/train/images",
             ...     annotations_directory_path=f"{dataset.location}/train/labels"
             ... )
 
-            >>> dataset.classes
+            >>> ds.classes
             ['dog', 'person']
             ```
         """
@@ -168,7 +234,7 @@ class Dataset:
         annotations = {
             image_name: detections for image_name, detections, _ in raw_annotations
         }
-        return Dataset(classes=classes, images=images, annotations=annotations)
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
 
     @classmethod
     def from_yolo(
@@ -177,7 +243,7 @@ class Dataset:
         annotations_directory_path: str,
         data_yaml_path: str,
         force_masks: bool = False,
-    ) -> Dataset:
+    ) -> DetectionDataset:
         """
         Creates a Dataset instance from YOLO formatted data.
 
@@ -188,7 +254,7 @@ class Dataset:
             force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
 
         Returns:
-            Dataset: A Dataset instance containing the loaded images and annotations.
+            DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
 
         Example:
             ```python
@@ -203,13 +269,13 @@ class Dataset:
             >>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
             >>> dataset = project.version(PROJECT_VERSION).download("yolov5")
 
-            >>> train_dataset = sv.Dataset.from_yolo(
+            >>> ds = sv.DetectionDataset.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
+            >>> ds.classes
             ['dog', 'person']
             ```
         """
@@ -219,4 +285,188 @@ class Dataset:
             data_yaml_path=data_yaml_path,
             force_masks=force_masks,
         )
-        return Dataset(classes=classes, images=images, annotations=annotations)
+        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 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)
+
+
+@dataclass
+class ClassificationDataset(BaseDataset):
+    """
+    Dataclass containing information about a classification 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, Classifications]
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def split(
+        self, split_ratio=0.8, random_state=None, shuffle: bool = True
+    ) -> Tuple[ClassificationDataset, ClassificationDataset]:
+        """
+        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.
+            random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
+            shuffle (bool, optional): Whether to shuffle the data before splitting.
+
+        Returns:
+            Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> cd = sv.ClassificationDataset(...)
+            >>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
+            >>> len(train_cd), len(test_cd)
+            (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 = ClassificationDataset(
+            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 = ClassificationDataset(
+            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_folder_structure(self, root_directory_path: str) -> None:
+        """
+        Saves the dataset as a multi-class folder structure.
+
+        Args:
+            root_directory_path (str): The path to the directory where the dataset will be saved.
+        """
+        os.makedirs(root_directory_path, exist_ok=True)
+
+        for class_name in self.classes:
+            os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
+
+        for image_name in self.images:
+            classification = self.annotations[image_name]
+            image = self.images[image_name]
+            class_id = (
+                classification.class_id[0]
+                if classification.confidence is None
+                else classification.get_top_k(1)[0]
+            )
+            class_name = self.classes[class_id]
+            image_path = os.path.join(root_directory_path, class_name, image_name)
+            cv2.imwrite(image_path, image)
+
+    @classmethod
+    def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
+        """
+        Load data from a multiclass folder structure into a ClassificationDataset.
+
+        Args:
+            root_directory_path (str): The path to the dataset directory.
+
+        Returns:
+            ClassificationDataset: The dataset.
+
+        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("folder")
+
+            >>> cd = sv.ClassificationDataset.from_folder_structure(
+            ...     root_directory_path=f"{dataset.location}/train"
+            ... )
+            ```
+        """
+        classes = os.listdir(root_directory_path)
+        classes = sorted(set(classes))
+
+        images = {}
+        annotations = {}
+
+        for class_name in classes:
+            class_id = classes.index(class_name)
+
+            for image in os.listdir(os.path.join(root_directory_path, class_name)):
+                image_dir = os.path.join(root_directory_path, class_name, image)
+                images[image] = cv2.imread(image_dir)
+                annotations[image] = Classifications(
+                    class_id=np.array([class_id]),
+                )
+
+        return cls(
+            classes=classes,
+            images=images,
+            annotations=annotations,
+        )

supervision/dataset/formats/pascal_voc.py

@@ -63,9 +63,6 @@ def detections_to_pascal_voc(
         str: An XML string in Pascal VOC format representing the detections.
     """
     height, width, depth = image_shape
-    image_area = height * width
-    minimum_detection_area = min_image_area_percentage * image_area
-    maximum_detection_area = max_image_area_percentage * image_area
 
     # Create root element
     annotation = Element("annotation")

supervision/dataset/formats/yolo.py

@@ -1,12 +1,19 @@
+import os
 from pathlib import Path
-from typing import Dict, List, Tuple, Union
+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
+from supervision.utils.file import (
+    list_files_with_extensions,
+    read_txt_file,
+    save_text_file,
+)
 
 
 def _parse_box(values: List[str]) -> np.ndarray:
@@ -78,6 +85,11 @@ def _extract_class_names(file_path: str) -> List[str]:
     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:
@@ -130,7 +142,7 @@ def load_yolo_annotations(
         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 paths as keys and images as values, and a dictionary with image paths as keys and corresponding Detections instances as values.
+        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"]
@@ -156,6 +168,93 @@ def load_yolo_annotations(
             lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
         )
 
-        images[str(image_path)] = image
-        annotations[str(image_path)] = annotation
+        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

@@ -1,5 +1,9 @@
-from typing import List
+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 (
@@ -8,6 +12,8 @@ from supervision.detection.utils import (
     mask_to_polygons,
 )
 
+T = TypeVar("T")
+
 
 def approximate_mask_with_polygons(
     mask: np.ndarray,
@@ -15,7 +21,7 @@ def approximate_mask_with_polygons(
     max_image_area_percentage: float = 1.0,
     approximation_percentage: float = 0.75,
 ) -> List[np.ndarray]:
-    height, width = mask
+    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
@@ -35,3 +41,41 @@ def approximate_mask_with_polygons(
         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/core.py

@@ -1,13 +1,19 @@
 from __future__ import annotations
 
 from dataclasses import astuple, dataclass
-from typing import Any, Iterator, List, Optional, Tuple
+from typing import Any, Iterator, List, Optional, Tuple, Union
 
+import cv2
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
+from supervision.detection.utils import (
+    extract_yolov8_masks,
+    non_max_suppression,
+    process_roboflow_result,
+    xywh_to_xyxy,
+)
 from supervision.geometry.core import Position
-from supervision.internal import deprecated
+from supervision.utils.internal import deprecated
 
 
 def _validate_xyxy(xyxy: Any, n: int) -> None:
@@ -191,6 +197,7 @@ class Detections:
             xyxy=yolov8_results.boxes.xyxy.cpu().numpy(),
             confidence=yolov8_results.boxes.conf.cpu().numpy(),
             class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
+            mask=extract_yolov8_masks(yolov8_results),
         )
 
     @classmethod
@@ -307,27 +314,14 @@ class Detections:
             >>> detections = sv.Detections.from_roboflow(roboflow_result, class_list)
             ```
         """
-        xyxy = []
-        confidence = []
-        class_id = []
-
-        for prediction in roboflow_result["predictions"]:
-            x = prediction["x"]
-            y = prediction["y"]
-            width = prediction["width"]
-            height = prediction["height"]
-            x_min = x - width / 2
-            y_min = y - height / 2
-            x_max = x_min + width
-            y_max = y_min + height
-            xyxy.append([x_min, y_min, x_max, y_max])
-            class_id.append(class_list.index(prediction["class"]))
-            confidence.append(prediction["confidence"])
-
+        xyxy, confidence, class_id, masks = process_roboflow_result(
+            roboflow_result=roboflow_result, class_list=class_list
+        )
         return Detections(
-            xyxy=np.array(xyxy),
-            confidence=np.array(confidence),
-            class_id=np.array(class_id).astype(int),
+            xyxy=xyxy,
+            confidence=confidence,
+            class_id=class_id,
+            mask=masks,
         )
 
     @classmethod
@@ -469,23 +463,43 @@ class Detections:
 
         raise ValueError(f"{anchor} is not supported.")
 
-    def __getitem__(self, index: np.ndarray) -> Detections:
-        if isinstance(index, np.ndarray) and (
-            index.dtype == bool or index.dtype == int
-        ):
-            return Detections(
-                xyxy=self.xyxy[index],
-                mask=self.mask[index] if self.mask is not None else None,
-                confidence=self.confidence[index]
-                if self.confidence is not None
-                else None,
-                class_id=self.class_id[index] if self.class_id is not None else None,
-                tracker_id=self.tracker_id[index]
-                if self.tracker_id is not None
-                else None,
-            )
-        raise TypeError(
-            f"Detections.__getitem__ not supported for index of type {type(index)}."
+    def __getitem__(
+        self, index: Union[int, slice, List[int], np.ndarray]
+    ) -> Detections:
+        """
+        Get a subset of the Detections object.
+
+        Args:
+            index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections
+
+        Returns:
+            (Detections): A subset of the Detections object.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> detections = sv.Detections(...)
+
+            >>> first_detection = detections[0]
+
+            >>> first_10_detections = detections[0:10]
+
+            >>> some_detections = detections[[0, 2, 4]]
+
+            >>> class_0_detections = detections[detections.class_id == 0]
+
+            >>> high_confidence_detections = detections[detections.confidence > 0.5]
+            ```
+        """
+        if isinstance(index, int):
+            index = [index]
+        return Detections(
+            xyxy=self.xyxy[index],
+            mask=self.mask[index] if self.mask is not None else None,
+            confidence=self.confidence[index] if self.confidence is not None else None,
+            class_id=self.class_id[index] if self.class_id is not None else None,
+            tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
         )
 
     @property

supervision/detection/line_counter.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional
 
 import cv2
 import numpy as np
@@ -6,7 +6,6 @@ import numpy as np
 from supervision.detection.core import Detections
 from supervision.draw.color import Color
 from supervision.geometry.core import Point, Rect, Vector
-from typing import Optional
 
 
 class LineZone:
@@ -71,6 +70,7 @@ class LineZone:
             else:
                 self.out_count += 1
 
+
 class LineZoneAnnotator:
     def __init__(
         self,
@@ -145,9 +145,16 @@ class LineZoneAnnotator:
             lineType=cv2.LINE_AA,
         )
 
-        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 = (
+            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

supervision/detection/tools/polygon_zone.py

@@ -17,7 +17,7 @@ class PolygonZone:
     A class for defining a polygon-shaped zone within a frame for detecting objects.
 
     Attributes:
-        polygon (np.ndarray): A numpy array defining the polygon vertices
+        polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points.
         frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
         triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
         current_count (int): The current count of detected objects within the zone
@@ -30,7 +30,7 @@ class PolygonZone:
         frame_resolution_wh: Tuple[int, int],
         triggering_position: Position = Position.BOTTOM_CENTER,
     ):
-        self.polygon = polygon
+        self.polygon = polygon.astype(int)
         self.frame_resolution_wh = frame_resolution_wh
         self.triggering_position = triggering_position
         self.current_count = 0

supervision/detection/utils.py

@@ -220,7 +220,9 @@ def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
     return np.array([x_min, y_min, x_max, y_max])
 
 
-def approximate_polygon(polygon: np.ndarray, percentage: float) -> np.ndarray:
+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.
 
@@ -228,11 +230,12 @@ def approximate_polygon(polygon: np.ndarray, percentage: float) -> np.ndarray:
     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).
+        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
+        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.
     """
 
@@ -246,8 +249,92 @@ def approximate_polygon(polygon: np.ndarray, percentage: float) -> np.ndarray:
 
     epsilon = 0
     approximated_points = polygon
-    while len(approximated_points) > target_points:
-        epsilon += 0.1
-        approximated_points = cv2.approxPolyDP(polygon, epsilon, closed=True)
+    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)
+
+
+def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]:
+    if not yolov8_results.masks:
+        return None
+
+    orig_shape = yolov8_results.orig_shape
+    inference_shape = tuple(yolov8_results.masks.data.shape[1:])
+
+    gain = 0
+    pad = (0, 0)
+
+    if inference_shape != orig_shape:
+        gain = min(
+            inference_shape[0] / orig_shape[0],
+            inference_shape[1] / orig_shape[1],
+        )
+        pad = (
+            (inference_shape[1] - orig_shape[1] * gain) / 2,
+            (inference_shape[0] - orig_shape[0] * gain) / 2,
+        )
+
+    top, left = int(pad[1]), int(pad[0])
+    bottom, right = int(inference_shape[0] - pad[1]), int(inference_shape[1] - pad[0])
+
+    mask_maps = []
+    masks = yolov8_results.masks.data.cpu().numpy()
+    for i in range(masks.shape[0]):
+        mask = masks[i]
+        mask = mask[top:bottom, left:right]
+        mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
+        mask_maps.append(mask)
+
+    return np.asarray(mask_maps, dtype=bool)
+
+
+def process_roboflow_result(
+    roboflow_result: dict, class_list: List[str]
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    if not roboflow_result["predictions"]:
+        return np.empty((0, 4)), np.empty(0), np.empty(0), None
+
+    xyxy = []
+    confidence = []
+    class_id = []
+    masks = []
+
+    image_width = int(roboflow_result["image"]["width"])
+    image_height = int(roboflow_result["image"]["height"])
+
+    for prediction in roboflow_result["predictions"]:
+        x = prediction["x"]
+        y = prediction["y"]
+        width = prediction["width"]
+        height = prediction["height"]
+        x_min = x - width / 2
+        y_min = y - height / 2
+        x_max = x_min + width
+        y_max = y_min + height
+
+        xyxy.append([x_min, y_min, x_max, y_max])
+        class_id.append(class_list.index(prediction["class"]))
+        confidence.append(prediction["confidence"])
+
+        if "points" not in prediction:
+            continue
+
+        polygon = np.array(
+            [[point["x"], point["y"]] for point in prediction["points"]], dtype=int
+        )
+
+        mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
+        masks.append(mask)
+
+    xyxy = np.array(xyxy)
+    confidence = np.array(confidence)
+    class_id = np.array(class_id).astype(int)
+    masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
+
+    return xyxy, confidence, class_id, masks

supervision/utils/file.py

@@ -17,13 +17,13 @@ def list_files_with_extensions(
 
     Examples:
         ```python
-        >>> from supervision import list_files_with_extensions
+        >>> import supervision as sv
 
         >>> # List all files in the directory
-        >>> files = list_files_with_extensions(directory='my_directory')
+        >>> files = sv.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'])
+        >>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
         ```
     """
     directory = Path(directory)
@@ -53,3 +53,16 @@ def read_txt_file(file_path: str) -> List[str]:
         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/utils/image.py

@@ -0,0 +1,93 @@
+import os
+import shutil
+from typing import Optional
+
+import cv2
+import numpy as np
+
+
+def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
+    """
+    Crops the given image based on the given bounding box.
+
+    Args:
+        image (np.ndarray): The image to be cropped, represented as a numpy array.
+        xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).
+
+    Returns:
+        (np.ndarray): The cropped image as a numpy array.
+
+    Examples:
+        ```python
+        >>> import supervision as sv
+
+        >>> detection = sv.Detections(...)
+        >>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
+        ...     for xyxy in detection.xyxy:
+        ...         cropped_image = sv.crop(image=image, xyxy=xyxy)
+        ...         sink.save_image(image=image)
+        ```
+    """
+
+    xyxy = np.round(xyxy).astype(int)
+    x1, y1, x2, y2 = xyxy
+    cropped_img = image[y1:y2, x1:x2]
+    return cropped_img
+
+
+class ImageSink:
+    def __init__(
+        self,
+        target_dir_path: str,
+        overwrite: bool = False,
+        image_name_pattern: str = "image_{:05d}.png",
+    ):
+        """
+        Initialize a context manager for saving images.
+
+        Args:
+            target_dir_path (str): The target directory where images will be saved.
+            overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
+            image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
+
+        Examples:
+            ```python
+            >>> import supervision as sv
+
+            >>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
+            ...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
+            ...         sink.save_image(image=image)
+            ```
+        """
+        self.target_dir_path = target_dir_path
+        self.overwrite = overwrite
+        self.image_name_pattern = image_name_pattern
+        self.image_count = 0
+
+    def __enter__(self):
+        if os.path.exists(self.target_dir_path):
+            if self.overwrite:
+                shutil.rmtree(self.target_dir_path)
+                os.makedirs(self.target_dir_path)
+        else:
+            os.makedirs(self.target_dir_path)
+
+        return self
+
+    def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
+        """
+        Save a given image in the target directory.
+
+        Args:
+            image (np.ndarray): The image to be saved.
+            image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
+        """
+        if image_name is None:
+            image_name = self.image_name_pattern.format(self.image_count)
+
+        image_path = os.path.join(self.target_dir_path, image_name)
+        cv2.imwrite(image_path, image)
+        self.image_count += 1
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        pass

supervision/utils/video.py

@@ -20,9 +20,9 @@ class VideoInfo:
 
     Examples:
         ```python
-        >>> from supervision import VideoInfo
+        >>> import supervision as sv
 
-        >>> video_info = VideoInfo.from_video_path(video_path='video.mp4')
+        >>> video_info = sv.VideoInfo.from_video_path(video_path='video.mp4')
 
         >>> video_info
         VideoInfo(width=3840, height=2160, fps=25, total_frames=538)
@@ -63,15 +63,15 @@ class VideoSink:
         target_path (str): The path to the output file where the video will be saved.
         video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
 
-    Examples:
+    Example:
         ```python
-        >>> from supervision import VideoInfo, VideoSink
+        >>> import supervision as sv
 
-        >>> video_info = VideoInfo.from_video_path(video_path='source_video.mp4')
+        >>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
 
-        >>> with VideoSink(target_path='target_video.mp4', video_info=video_info) as s:
-        ...     frame = ...
-        ...     s.write_frame(frame=frame)
+        >>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
+        ...     for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
+        ...         sink.write_frame(frame=frame)
         ```
     """
 
@@ -93,35 +93,43 @@ class VideoSink:
     def write_frame(self, frame: np.ndarray):
         self.__writer.write(frame)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_value, exc_traceback):
         self.__writer.release()
 
 
-def get_video_frames_generator(source_path: str) -> Generator[np.ndarray, None, None]:
+def get_video_frames_generator(
+    source_path: str, stride: int = 1
+) -> Generator[np.ndarray, None, None]:
     """
     Get a generator that yields the frames of the video.
 
     Args:
         source_path (str): The path of the video file.
+        stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each.
 
     Returns:
         (Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
 
     Examples:
         ```python
-        >>> from supervision import get_video_frames_generator
+        >>> import supervision as sv
 
-        >>> for frame in get_video_frames_generator(source_path='source_video.mp4'):
+        >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
         ...     ...
         ```
     """
     video = cv2.VideoCapture(source_path)
     if not video.isOpened():
         raise Exception(f"Could not open video at {source_path}")
+
+    frame_count = 0
     success, frame = video.read()
     while success:
-        yield frame
+        if frame_count % stride == 0:
+            yield frame
         success, frame = video.read()
+        frame_count += 1
+
     video.release()
 
 

test/classification/test_core.py

@@ -0,0 +1,63 @@
+from typing import List, Optional, Tuple
+from contextlib import ExitStack as DoesNotRaise
+
+import pytest
+import numpy as np
+
+from supervision.classification.core import Classifications
+
+
+@pytest.mark.parametrize(
+    "class_id, confidence, k, expected_result, exception",
+    [
+        (
+            np.array([0, 1, 2, 3, 4]),
+            np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
+            5,
+            (np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
+            DoesNotRaise(),
+        ),  # class_id with 5 numbers and 5 confidences
+        (
+            np.array([5, 1, 2, 3, 4]),
+            np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
+            1,
+            (np.array([2]), np.array([0.9])),
+            DoesNotRaise(),
+        ),  # class_id with 5 numbers and 5 confidences, retrieve where k = 1
+        (
+            np.array([4, 1, 2, 3, 6, 5]),
+            np.array([0.8, 0.2, 0.9, 0.4, 0.5, 0.1]),
+            2,
+            (np.array([2, 4]), np.array([0.9, 0.8])),
+            DoesNotRaise(),
+        ),  # class_id with 5 numbers and 5 confidences, retrieve where k = 3
+        (
+            np.array([0, 1, 2, 3, 4]),
+            np.array([]),
+            5,
+            None,
+            pytest.raises(ValueError),
+        ),  # class_id with 5 numbers and 0 confidences
+        (
+            [0, 1, 2, 3, 4],
+            [0.1, 0.2, 0.3, 0.4],
+            5,
+            None,
+            pytest.raises(ValueError),
+        ),  # class_id with 5 numbers and 4 confidences
+    ],
+)
+def test_top_k(
+    class_id: np.ndarray,
+    confidence: Optional[np.ndarray],
+    k: int,
+    expected_result: Optional[Tuple[np.ndarray, np.ndarray]],
+    exception: Exception,
+) -> None:
+    with exception:
+        result = Classifications(
+            class_id=np.array(class_id), confidence=np.array(confidence)
+        ).get_top_k(k)
+
+        assert np.array_equal(result[0], expected_result[0])
+        assert np.array_equal(result[1], expected_result[1])

test/dataset/formats/test_yolo.py

@@ -5,7 +5,8 @@ import pytest
 import numpy as np
 
 from supervision.detection.core import Detections
-from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask
+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:
@@ -202,3 +203,102 @@ def test_yolo_annotations_to_detections(
         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

@@ -28,13 +28,6 @@ DETECTIONS = Detections(
 )
 
 
-PREDICTIONS = np.array([
-        [       2254,         906,        2447,        1353,     0.90538,           0],
-        [       2049,        1133,        2226,        1371,     0.59002,          56],
-        [        727,        1224,         838,        1601,     0.51119,          39],
-    ], dtype=np.float32)
-
-
 def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections:
     return Detections(
         xyxy = np.array(xyxy, dtype=np.float32),
@@ -101,11 +94,108 @@ def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None)
             ),
             DoesNotRaise()
         ),  # take no detections
+        (
+            DETECTIONS,
+            [0, 2],
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353],
+                    [        727,        1224,         838,        1601]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538,
+                    0.51119
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0,
+                    39
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only first and third detection using List[int] index
+        (
+            DETECTIONS,
+            np.array([0, 2]),
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353],
+                    [        727,        1224,         838,        1601]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538,
+                    0.51119
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0,
+                    39
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only first and third detection using np.ndarray index
+        (
+            DETECTIONS,
+            0,
+            Detections(
+                xyxy=np.array([
+                    [       2254,         906,        2447,        1353]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.90538
+                ], dtype=np.float32),
+                class_id=np.array([
+                    0
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only first detection by index
+        (
+            DETECTIONS,
+            slice(1, 3),
+            Detections(
+                xyxy=np.array([
+                    [       2049,        1133,        2226,        1371],
+                    [        727,        1224,         838,        1601]
+                ], dtype=np.float32),
+                confidence=np.array([
+                    0.59002,
+                    0.51119
+                ], dtype=np.float32),
+                class_id=np.array([
+                    56,
+                    39
+                ], dtype=int)
+            ),
+            DoesNotRaise()
+        ),  # take only first detection by index slice (1, 3)
+        (
+            DETECTIONS,
+            10,
+            None,
+            pytest.raises(IndexError)
+        ),  # index out of range
+        (
+            DETECTIONS,
+            [0, 2, 10],
+            None,
+            pytest.raises(IndexError)
+        ),  # index out of range
+        (
+            DETECTIONS,
+            np.array([0, 2, 10]),
+            None,
+            pytest.raises(IndexError)
+        ),
+        (
+            DETECTIONS,
+            np.array([True, True, True, True, True, True, True, True, True, True, True]),
+            None,
+            pytest.raises(IndexError)
+        )
     ]
 )
 def test_getitem(
         detections: Detections,
-        index: Union[int, slice, np.ndarray],
+        index:  Union[int, slice, List[int], np.ndarray],
         expected_result: Optional[Detections],
         exception: Exception
 ) -> None:

test/detection/test_utils.py

@@ -5,7 +5,8 @@ import pytest
 
 import numpy as np
 
-from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
+from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area, \
+    process_roboflow_result
 
 
 @pytest.mark.parametrize(
@@ -276,3 +277,61 @@ def test_filter_polygons_by_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)
+
+
+@pytest.mark.parametrize(
+    "roboflow_result, class_list, expected_result, exception",
+    [
+        (
+            {
+                "predictions": [],
+                "image": {"width": 1000, "height": 1000}
+            },
+            ["person", "car", "truck"],
+            (
+                np.empty((0, 4)),
+                np.empty(0),
+                np.empty(0),
+                None
+             ),
+            DoesNotRaise()
+        ),  # empty result
+        (
+            {
+                "predictions": [
+                    {
+                        "x": 200.0,
+                        "y": 300.0,
+                        "width": 50.0,
+                        "height": 50.0,
+                        "confidence": 0.9,
+                        "class": "person"
+                    }
+                ],
+                "image": {"width": 1000, "height": 1000}
+            },
+            ["person", "car", "truck"],
+            (
+                np.array([
+                    [175.0, 275.0, 225.0, 325.0]
+                ]),
+                np.array([0.9]),
+                np.array([0]),
+                None
+             ),
+            DoesNotRaise()
+        ),  # single bounding box
+    ]
+)
+def test_process_roboflow_result(
+        roboflow_result: dict,
+        class_list: List[str],
+        expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
+        exception: Exception
+) -> None:
+    with exception:
+        result = process_roboflow_result(roboflow_result=roboflow_result, class_list=class_list)
+        assert np.array_equal(result[0], expected_result[0])
+        assert np.array_equal(result[1], expected_result[1])
+        assert np.array_equal(result[2], expected_result[2])
+        assert (result[3] is None and expected_result[3] is None) or (np.array_equal(result[3], expected_result[3]))