🛡️ add badges to README.md

feature/dataset_split

README.md

@@ -56,6 +56,10 @@
 
   <br>
 
+[![version](https://badge.fury.io/py/supervision.svg)](https://badge.fury.io/py/supervision)
+[![downloads](https://img.shields.io/pypi/dm/supervision)](https://pypistats.org/packages/supervision)
+[![license](https://img.shields.io/pypi/l/supervision)](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
+[![python-version](https://img.shields.io/pypi/pyversions/supervision)](https://badge.fury.io/py/supervision)
 [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
 
 </div>

docs/changelog.md

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

docs/dataset/core.md

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

setup.py

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

supervision/__init__.py

@@ -1,6 +1,6 @@
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 
-from supervision.dataset.core import Dataset
+from supervision.dataset.core import BaseDataset, DetectionDataset
 from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
 from supervision.detection.core import Detections
 from supervision.detection.line_counter import LineZone, LineZoneAnnotator

supervision/dataset/core.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+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
@@ -11,15 +12,33 @@ 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
 
 
 @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 +70,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. Default is 0.8.
+            random_state (int, optional): The seed for the random number generator. This is used for reproducibility. Default is None.
+            shuffle (bool, optional): Whether to shuffle the data before splitting. Default is True.
+
+        Returns:
+            Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> dataset = sv.DetectionDataset(...)
+            >>> train_dataset, test_dataset = dataset.split(split_ratio=0.7, random_state=42, shuffle=True)
+            >>> len(train_dataset), len(test_dataset)
+            (700, 300)
+            ```
+        """
+
+        image_names = list(self.images.keys())
+        train_names, test_names = train_test_split(
+            data=image_names,
+            train_ratio=split_ratio,
+            random_state=random_state,
+            shuffle=shuffle,
+        )
+
+        train_dataset = DetectionDataset(
+            classes=self.classes,
+            images={name: self.images[name] for name in train_names},
+            annotations={name: self.annotations[name] for name in train_names},
+        )
+        test_dataset = DetectionDataset(
+            classes=self.classes,
+            images={name: self.images[name] for name in test_names},
+            annotations={name: self.annotations[name] for name in test_names},
+        )
+        return train_dataset, test_dataset
+
     def as_pascal_voc(
         self,
         images_directory_path: Optional[str] = None,
         annotations_directory_path: Optional[str] = None,
         min_image_area_percentage: float = 0.0,
         max_image_area_percentage: float = 1.0,
-        approximation_percentage: float = 0.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 +171,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 +180,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
@@ -168,7 +232,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 +241,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 +252,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
@@ -219,4 +283,50 @@ 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 (You Only Look Once) format. This method saves the images and their corresponding
+        annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
+        for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
+
+        The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
+
+        Args:
+            images_directory_path (Optional[str]): The path to the directory where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path (Optional[str]): The path to the directory where the annotations in
+                YOLO format should be saved. If not provided, annotations will not be saved.
+            data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
+                If not provided, the file will not be saved.
+            min_image_area_percentage (float): The minimum percentage of detection area relative to
+                the image area for a detection to be included.
+            max_image_area_percentage (float): The maximum percentage of detection area relative to
+                the image area for a detection to be included.
+            approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
+                in the range [0, 1). This is useful for simplifying the annotations.
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                images_directory_path=images_directory_path, images=self.images
+            )
+        if annotations_directory_path is not None:
+            save_yolo_annotations(
+                annotations_directory_path=annotations_directory_path,
+                images=self.images,
+                annotations=self.annotations,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+            )
+        if data_yaml_path is not None:
+            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)

supervision/dataset/formats/pascal_voc.py

@@ -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,15 @@
+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.file import list_files_with_extensions, read_txt_file, save_text_file
 
 
 def _parse_box(values: List[str]) -> np.ndarray:
@@ -78,6 +81,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 +138,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 +164,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/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/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,12 @@ 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)

supervision/file.py

@@ -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")

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