Merge pull request #136 from roboflow/0.10.0-release

`0.10.0` release
🪵 changelog.md updates before 0.10.0 release
2023-06-14 16:50:24 +02:00 · 2023-06-14 16:49:10 +02:00 · 2023-06-14 16:40:49 +02:00 · 2023-06-14 16:38:08 +02:00 · 2023-06-14 15:54:40 +02:00 · 2023-06-14 15:52:20 +02:00
34 changed files with 2436 additions and 191 deletions
@@ -0,0 +1 @@
+*.ipynb linguist-vendored
@@ -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
@@ -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,77 +281,4 @@
      </a>
  </div>

-  <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>
-
-## 👋 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!
@@ -1,8 +1,98 @@
+### 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. 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`.
+- 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>
@@ -0,0 +1,3 @@
+## Classifications
+
+:::supervision.classification.core.Classifications
@@ -5,4 +5,8 @@

 ## DetectionDataset

-:::supervision.dataset.core.DetectionDataset
+:::supervision.dataset.core.DetectionDataset
+
+## ClassificationDataset
+
+:::supervision.dataset.core.ClassificationDataset
@@ -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)"

@@ -1,7 +0,0 @@
-## plot_image
-
-:::supervision.notebook.utils.plot_image
-
-## plot_images_grid
-
-:::supervision.notebook.utils.plot_images_grid
@@ -0,0 +1,3 @@
+## list_files_with_extensions
+
+:::supervision.utils.file.list_files_with_extensions
@@ -0,0 +1,7 @@
+## ImageSink
+
+:::supervision.utils.image.ImageSink
+
+## crop
+
+:::supervision.utils.image.crop
@@ -0,0 +1,7 @@
+## plot_image
+
+:::supervision.utils.notebook.plot_image
+
+## plot_images_grid
+
+:::supervision.utils.notebook.plot_images_grid
@@ -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
@@ -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
@@ -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:
@@ -52,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',
@@ -63,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',
 )
@@ -1,6 +1,11 @@
-__version__ = "0.8.0"
+__version__ = "0.10.0"

-from supervision.dataset.core import BaseDataset, DetectionDataset
+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,
@@ -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
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import os
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
@@ -8,6 +9,7 @@ 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,
@@ -19,7 +21,7 @@ from supervision.dataset.formats.yolo import (
 )
 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
@@ -77,9 +79,9 @@ class DetectionDataset(BaseDataset):
        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.
+            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.
@@ -88,9 +90,9 @@ class DetectionDataset(BaseDataset):
            ```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)
+            >>> 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)
            ```
        """
@@ -195,12 +197,12 @@ class DetectionDataset(BaseDataset):
            >>> 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']
            ```
        """
@@ -267,13 +269,13 @@ class DetectionDataset(BaseDataset):
            >>> 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']
            ```
        """
@@ -295,7 +297,7 @@ class DetectionDataset(BaseDataset):
        approximation_percentage: float = 0.0,
    ) -> None:
        """
-        Exports the dataset to YOLO (You Only Look Once) format. This method saves the images and their corresponding
+        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.

@@ -330,3 +332,141 @@ class DetectionDataset(BaseDataset):
            )
        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,
+        )
@@ -9,7 +9,11 @@ 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
+from supervision.utils.file import (
+    list_files_with_extensions,
+    read_txt_file,
+    save_text_file,
+)


 def _parse_box(values: List[str]) -> np.ndarray:
@@ -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
@@ -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
@@ -258,3 +258,83 @@ def approximate_polygon(
            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
@@ -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)
@@ -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
@@ -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()


@@ -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])
@@ -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:
@@ -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]))
Author	SHA1	Message	Date
Piotr Skalski	9cff624e15	Merge pull request #136 from roboflow/0.10.0-release `0.10.0` release	2023-06-14 16:50:24 +02:00
SkalskiP	584c4fef82	🪵 changelog.md updates before `0.10.0` release	2023-06-14 16:49:10 +02:00
Piotr Skalski	282aa2c5d3	Merge pull request #110 from roboflow/rf-polygon-loading Add semantic segmentation polygon loading from Roboflow	2023-06-14 16:40:49 +02:00
SkalskiP	453d4df068	🧪 refactor + some tests	2023-06-14 16:38:08 +02:00
SkalskiP	4d4b38b292	🖤 make black happy	2023-06-14 15:54:40 +02:00
SkalskiP	13e8afa996	👌 final changes before `0.10.0` release	2023-06-14 15:52:20 +02:00
James Gallagher	dd4cae328b	add missing import from merge, run linters	2023-06-14 14:21:15 +01:00
James	ed7103851a	Merge branch 'main' into rf-polygon-loading	2023-06-14 14:14:03 +01:00
James Gallagher	ca40b001dc	fix bug in mask loading	2023-06-14 14:08:53 +01:00
Piotr Skalski	de3ccc5c5d	Merge pull request #125 from roboflow/add-classification Add ClassificationDataset API	2023-06-14 14:11:27 +02:00
SkalskiP	1b6d02898c	8️⃣ `Classifications.from_yolov8` added	2023-06-14 14:07:14 +02:00
SkalskiP	5ac92bc43b	🧪 ready for final tests	2023-06-14 13:15:28 +02:00
Piotr Skalski	86d9293770	Merge pull request #135 from dankresio/patch-1 Update stride definition	2023-06-14 11:08:11 +02:00
Daniel M. García-Ocaña	0e943e673e	Update stride definition	2023-06-14 10:53:06 +02:00
James Gallagher	234e376fca	respond to feedback	2023-06-14 09:48:00 +01:00
SkalskiP	54b7cbeb70	Merge remote-tracking branch 'origin/main' # Conflicts: # README.md	2023-06-13 17:47:25 +02:00
SkalskiP	d000292eb2	✨ new demo notebook	2023-06-13 17:45:22 +02:00
SkalskiP	b551e6d410	✨ new demo notebook	2023-06-13 17:34:37 +02:00
James Gallagher	138c271af0	update as_multiclass logic	2023-06-13 14:39:32 +01:00
James Gallagher	d22940ce30	add classifications to table of contents	2023-06-13 14:36:31 +01:00
James Gallagher	06da0e70d2	remove unused files	2023-06-13 14:35:52 +01:00
James Gallagher	5d8b378f35	add test cases for get_top_k, add docs for Classifications and ClassificationsDataset	2023-06-13 14:35:32 +01:00
James Gallagher	a4830389de	fix import	2023-06-13 13:24:11 +01:00
James Gallagher	d9a2efc64b	remove unused imports, fix top k bug, fix import paths	2023-06-13 13:16:48 +01:00
James Gallagher	f91b844b7d	fix ClassificationDataset import	2023-06-13 13:12:04 +01:00
James Gallagher	dc45d47e7e	run linter	2023-06-13 12:48:34 +01:00
James Gallagher	b5dedf09cd	refactor code, respond to feedback	2023-06-13 12:48:07 +01:00
James Gallagher	eb1ae1b790	respond to feedback	2023-06-13 12:00:46 +01:00
James Gallagher	f2f8a564e0	run linter	2023-06-13 11:38:19 +01:00
James Gallagher	140a219091	generalise as_multiclass_folder_structure	2023-06-13 11:38:02 +01:00
James Gallagher	86af9b9dce	continue work on classifications	2023-06-13 11:13:38 +01:00
Piotr Skalski	4366e4480f	`💜Built with supervision` added to README.md	2023-06-12 19:06:08 +02:00
James Gallagher	b66e9c8c32	change val to valid, remove example code	2023-06-08 11:51:41 +01:00
James Gallagher	d0d5626156	run linter	2023-06-07 16:29:52 +01:00
James Gallagher	c661bf250d	add classification dataset modules	2023-06-07 16:25:22 +01:00
SkalskiP	ccb32f8aa1	🪵 Update changelog.	2023-06-07 12:49:53 +02:00
Piotr Skalski	97b9edcbec	Merge pull request #122 from roboflow/feature/75_crop_bounding_boxes_as_separate_image feature/75_crop_bounding_boxes_as_separate_image	2023-06-07 12:20:00 +02:00
SkalskiP	6b52863c9f	📚 Documentation update	2023-06-07 12:15:14 +02:00
SkalskiP	2fcc035413	🖼️ Add utility to crop image based on xyxy box	2023-06-07 12:12:39 +02:00
Piotr Skalski	7c8f16bf0f	Merge pull request #121 from roboflow/fix/106_better_handling_of_wrong_data_types_of_mask 🛠️ Fix #106 - now you can pass polygon with int and float vertices.	2023-06-07 11:58:43 +02:00
SkalskiP	1870670c77	🛠️ Fix #106 - now you can pass polygon with int and float vertices.	2023-06-07 11:54:50 +02:00
SkalskiP	802347f4d8	🧹 cleanup before `supervision` `0.9.0` release	2023-06-07 11:47:34 +02:00
Piotr Skalski	920019c887	Merge pull request #120 from roboflow/feature/new_video_utils feature/new_video_utils	2023-06-07 00:26:05 +02:00
SkalskiP	477b9a8f38	🛠️ fix `ImageSink.__enter__`	2023-06-06 21:58:26 +02:00
SkalskiP	14f62deacf	📚 `ImageSink` documentation updated	2023-06-06 20:13:33 +02:00
SkalskiP	08eb513594	📚 `ImageSink` documentation added	2023-06-06 19:58:30 +02:00
SkalskiP	a1615fd275	🎬 small `VideoSink` documentation update. + 🖼️ `ImageSink` added	2023-06-06 19:54:18 +02:00
SkalskiP	ce979216f2	🐍 small `setup.py` and `README.md` updates	2023-06-06 18:01:06 +02:00
SkalskiP	95f94319f8	😷 Move YOLOv8 mask processing code to separate `extract_yolov8_masks` util	2023-06-06 17:50:14 +02:00
SkalskiP	98f60c24d3	😷 Add `stride` to `get_video_frames_generator` function. Initial `from_yolov8` with `masks` refactor.	2023-06-06 17:35:40 +02:00
Piotr Skalski	1c043d6cc7	Merge pull request #101 from hardikdava/main Adding yolov8 masks to detections	2023-06-06 10:52:40 +02:00
Piotr Skalski	86aee5ed98	Merge pull request #118 from roboflow/feature/115_how_to_select_detections_by_index feature/115_how_to_select_detections_by_index	2023-06-05 11:34:35 +02:00
SkalskiP	0d38b95da7	📚update documentation + 🖤 make black happy	2023-06-05 11:31:52 +02:00
SkalskiP	97892e3138	⛏️pick `Detections` by index	2023-06-05 11:01:14 +02:00
Piotr Skalski	17a42bc728	Merge pull request #117 from lomnes-atlast-food/update-readme-loop-dataset Update README.md example for looping through dataset	2023-06-02 08:25:17 +02:00
Stephen Lomnes	be65ba7370	Update README.md Example in README documentation out of date for v.0.8.0. __iter__() yields a length 3 tuple. Example shows a length 2 tuple. Change proposed: add image_name to example. Reference: From API docs ``` __iter__() ¶ Iterate over the images and annotations in the dataset. Yields: Type Description Iterator[Tuple[str, np.ndarray, Detections]] Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name, the image data, and its corresponding annotation. ```	2023-06-01 09:08:33 -04:00
Hardik Dava	72d86ba180	Fixing for uneven image size	2023-05-24 17:27:35 +02:00
Hardik Dava	fba83ec371	Fixed minor typo bug	2023-05-24 17:22:16 +02:00
Hardik Dava	4a81a17ddb	yolov8 masks plot using with letterbox concept - gain and pads are calculated using yolov8 results. - resize masks to input image size	2023-05-24 15:02:00 +02:00
Hardik Dava	30d253d5cc	Fixed import error	2023-05-23 18:40:48 +02:00
James Gallagher	544f9cb0e2	incorporate feedback	2023-05-23 13:02:58 +01:00
James Gallagher	84928a2da9	add rf segmentation polygon loading	2023-05-23 12:02:51 +01:00
Hardik Dava	08ae0b2fc7	using basic functions from cv2	2023-05-20 19:22:30 +02:00
Piotr Skalski	0261bec2ce	Merge pull request #104 from roboflow/SkalskiP-patch-1 📚 Update README.md	2023-05-18 01:10:05 +02:00
Piotr Skalski	5a4cb1b302	Update README.md	2023-05-18 01:08:49 +02:00
Piotr Skalski	91655f9609	🖌️ final README.md updates	2023-05-18 01:07:36 +02:00
Piotr Skalski	42d4079b23	🎬 video clap section	2023-05-18 01:04:12 +02:00
Piotr Skalski	f03a47a5df	🔥 quick start section	2023-05-18 00:59:01 +02:00
Piotr Skalski	c01efeaaca	📚 Update README.md	2023-05-17 23:08:28 +02:00
SkalskiP	5489d405c6	📚 `DetectionDataset` documentation update	2023-05-17 21:46:45 +02:00
SkalskiP	5f97eac9a9	📚 `DetectionDataset` documentation update	2023-05-17 21:35:33 +02:00
Hardik Dava	b474a47065	Adding yolov8 masks to detections	2023-05-17 18:38:40 +02:00