superiorvision 0.30.0.dev0__py3-none-any.whl

supervision/dataset/core.py

@@ -0,0 +1,1313 @@
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from copy import deepcopy
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from typing import cast
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from tqdm.auto import tqdm
+
+from supervision.classification.core import Classifications
+from supervision.config import CLASS_NAME_DATA_FIELD
+from supervision.dataset.formats.coco import (
+    load_coco_annotations,
+    save_coco_annotations,
+)
+from supervision.dataset.formats.createml import (
+    load_createml_annotations,
+    save_createml_annotations,
+)
+from supervision.dataset.formats.labelme import (
+    load_labelme_annotations,
+    save_labelme_annotations,
+)
+from supervision.dataset.formats.pascal_voc import (
+    load_pascal_voc_annotations,
+    save_pascal_voc_annotations,
+)
+from supervision.dataset.formats.yolo import (
+    load_yolo_annotations,
+    save_data_yaml,
+    save_yolo_annotations,
+)
+from supervision.dataset.utils import (
+    build_class_index_mapping,
+    check_no_basename_collisions,
+    map_detections_class_id,
+    merge_class_lists,
+    save_dataset_images,
+    train_test_split,
+)
+from supervision.detection.core import Detections
+from supervision.utils.internal import warn_deprecated
+from supervision.utils.iterables import find_duplicates
+
+_IMAGE_FILE_EXTENSIONS = frozenset(
+    {".bmp", ".jpeg", ".jpg", ".png", ".tif", ".tiff", ".webp"}
+)
+
+
+class BaseDataset(ABC):
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
+
+    @abstractmethod
+    def split(
+        self,
+        split_ratio: float = 0.8,
+        random_state: int | None = None,
+        shuffle: bool = True,
+    ) -> tuple[BaseDataset, BaseDataset]:
+        pass
+
+
+class DetectionDataset(BaseDataset):
+    """
+    Contains information about a detection dataset. Handles lazy image loading
+    and annotation retrieval, dataset splitting, conversions into multiple
+    formats.
+
+    Attributes:
+        classes: List containing dataset class names.
+        images:
+            Accepts a list of image paths. Passing a dict
+            (``Dict[str, np.ndarray]``) is deprecated in ``0.30.0`` and will
+            be removed in ``0.33.0``; use a list of paths instead.
+            When a list of paths is provided, images are loaded lazily on
+            demand, which is more memory-efficient.
+        annotations: Dictionary mapping
+            image path to annotations. The dictionary keys match
+            match the keys in `images` or entries in the list of
+            image paths.
+    """
+
+    def __init__(
+        self,
+        classes: list[str],
+        images: list[str] | dict[str, npt.NDArray[np.uint8]],
+        annotations: dict[str, Detections],
+    ) -> None:
+        self.classes = classes
+
+        if set(images) != set(annotations):
+            raise ValueError(
+                "The keys of the images and annotations dictionaries must match."
+            )
+        self.annotations = {
+            image_path: deepcopy(annotation)
+            for image_path, annotation in annotations.items()
+        }
+
+        np_classes = np.array(self.classes)
+        for image_path, annotation in self.annotations.items():
+            class_ids = annotation.class_id
+            if class_ids is None:
+                continue
+            if not np.issubdtype(class_ids.dtype, np.integer):
+                raise ValueError(
+                    f"Detection annotation for image {image_path!r} contains "
+                    f"non-integer class_id values with dtype {class_ids.dtype}."
+                )
+
+            invalid_class_ids = class_ids[
+                (class_ids < 0) | (class_ids >= len(self.classes))
+            ]
+            if len(invalid_class_ids) > 0:
+                valid_range = (
+                    "empty"
+                    if len(self.classes) == 0
+                    else f"[0, {len(self.classes) - 1}]"
+                )
+                raise ValueError(
+                    f"Detection annotation for image {image_path!r} contains "
+                    f"class_id {int(invalid_class_ids[0])}, outside the valid "
+                    f"range {valid_range} for {len(self.classes)} classes."
+                )
+
+            annotation.data[CLASS_NAME_DATA_FIELD] = np_classes[class_ids]
+
+        # Eliminate duplicates while preserving order
+        self.image_paths = list(dict.fromkeys(images))
+
+        self._images_in_memory: dict[str, npt.NDArray[np.uint8]] = {}
+        if isinstance(images, dict):
+            self._images_in_memory = images
+            warn_deprecated(
+                "Passing a `Dict[str, np.ndarray]` into `DetectionDataset` is "
+                "deprecated in `0.30.0` and will be removed in `0.33.0`. Use "
+                "a list of paths `List[str]` instead."
+            )
+
+    def _get_image(self, image_path: str) -> npt.NDArray[np.uint8]:
+        """Assumes that image is in dataset."""
+        if self._images_in_memory:
+            return self._images_in_memory[image_path]
+        image = cv2.imread(image_path)
+        if image is None:
+            raise ValueError(f"Could not read image from path: {image_path}")
+        return cast(npt.NDArray[np.uint8], image)
+
+    def __len__(self) -> int:
+        return len(self._images_in_memory) or len(self.image_paths)
+
+    def __getitem__(self, i: int) -> tuple[str, npt.NDArray[np.uint8], Detections]:
+        """
+        Returns:
+            The image path, image data,
+                and its corresponding annotation at index i.
+        """
+        image_path = self.image_paths[i]
+        image = self._get_image(image_path)
+        annotation = self.annotations[image_path]
+        return image_path, image, annotation
+
+    def __iter__(self) -> Iterator[tuple[str, npt.NDArray[np.uint8], Detections]]:
+        """
+        Iterate over the images and annotations in the dataset.
+
+        Yields:
+            Tuples containing the image path, image data, and its annotation.
+        """
+        for i in range(len(self)):
+            image_path, image, annotation = self[i]
+            yield image_path, image, annotation
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, DetectionDataset):
+            return False
+
+        if self.classes != other.classes:
+            return False
+
+        if self.image_paths != other.image_paths:
+            return False
+
+        if self._images_in_memory or other._images_in_memory:
+            if not np.array_equal(
+                list(self._images_in_memory.values()),
+                list(other._images_in_memory.values()),
+            ):
+                return False
+
+        if self.annotations != other.annotations:
+            return False
+
+        return True
+
+    def split(
+        self,
+        split_ratio: float = 0.8,
+        random_state: int | None = None,
+        shuffle: bool = True,
+    ) -> tuple[DetectionDataset, DetectionDataset]:
+        """
+        Splits the dataset into two parts (training and testing)
+            using the provided split_ratio. The input dataset is not mutated.
+
+        Args:
+            split_ratio: The ratio of the training
+                set to the entire dataset.
+            random_state: The seed for the random number generator.
+                This is used for reproducibility.
+            shuffle: Whether to shuffle the data before splitting.
+
+        Returns:
+            A tuple containing
+                the training and testing datasets.
+
+        Examples:
+            ```pycon
+            >>> import numpy as np
+            >>> import supervision as sv
+            >>> ds = sv.DetectionDataset(
+            ...     classes=['dog', 'person'],
+            ...     images={
+            ...         'img1.jpg': np.zeros((100, 100, 3), dtype=np.uint8),
+            ...         'img2.jpg': np.zeros((100, 100, 3), dtype=np.uint8),
+            ...     },
+            ...     annotations={
+            ...         'img1.jpg': sv.Detections(xyxy=np.array([[10, 10, 20, 20]])),
+            ...         'img2.jpg': sv.Detections(xyxy=np.array([[30, 30, 40, 40]])),
+            ...     }
+            ... )
+            >>> train_ds, test_ds = ds.split(split_ratio=0.5, random_state=42)
+            >>> len(train_ds), len(test_ds)
+            (1, 1)
+
+            ```
+        """
+
+        train_paths, test_paths = train_test_split(
+            data=self.image_paths,
+            train_ratio=split_ratio,
+            random_state=random_state,
+            shuffle=shuffle,
+        )
+
+        train_annotations = {path: self.annotations[path] for path in train_paths}
+        test_annotations = {path: self.annotations[path] for path in test_paths}
+
+        train_dataset = DetectionDataset(
+            classes=self.classes,
+            images=train_paths,
+            annotations=train_annotations,
+        )
+        test_dataset = DetectionDataset(
+            classes=self.classes,
+            images=test_paths,
+            annotations=test_annotations,
+        )
+        if self._images_in_memory:
+            train_dataset._images_in_memory = {
+                path: self._images_in_memory[path] for path in train_paths
+            }
+            test_dataset._images_in_memory = {
+                path: self._images_in_memory[path] for path in test_paths
+            }
+        return train_dataset, test_dataset
+
+    @classmethod
+    def merge(cls, dataset_list: list[DetectionDataset]) -> DetectionDataset:
+        """
+        Merge a list of `DetectionDataset` objects into a single
+            `DetectionDataset` object.
+
+        This method takes a list of `DetectionDataset` objects and combines
+        their respective fields (`classes`, `images`,
+        `annotations`) into a single `DetectionDataset` object.
+
+        Args:
+            dataset_list: A list of `DetectionDataset`
+                objects to merge.
+
+        Returns:
+            A single `DetectionDataset` object containing
+            the merged data from the input list.
+
+        Examples:
+            ```pycon
+            >>> import numpy as np
+            >>> import supervision as sv
+            >>> ds_1 = sv.DetectionDataset(
+            ...     classes=['dog', 'person'],
+            ...     images={'img1.jpg': np.zeros((100, 100, 3), dtype=np.uint8)},
+            ...     annotations={'img1.jpg': sv.Detections.empty()}
+            ... )
+            >>> len(ds_1)
+            1
+            >>> ds_1.classes
+            ['dog', 'person']
+            >>> ds_2 = sv.DetectionDataset(
+            ...     classes=['cat'],
+            ...     images={'img2.jpg': np.zeros((100, 100, 3), dtype=np.uint8)},
+            ...     annotations={'img2.jpg': sv.Detections.empty()}
+            ... )
+            >>> len(ds_2)
+            1
+            >>> ds_2.classes
+            ['cat']
+            >>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
+            >>> len(ds_merged)
+            2
+            >>> ds_merged.classes
+            ['cat', 'dog', 'person']
+
+            ```
+        """
+
+        def is_in_memory(dataset: DetectionDataset) -> bool:
+            return len(dataset._images_in_memory) > 0 or len(dataset.image_paths) == 0
+
+        def is_lazy(dataset: DetectionDataset) -> bool:
+            return len(dataset._images_in_memory) == 0
+
+        all_in_memory = all([is_in_memory(dataset) for dataset in dataset_list])
+        all_lazy = all([is_lazy(dataset) for dataset in dataset_list])
+        if not all_in_memory and not all_lazy:
+            raise ValueError(
+                "Merging lazy and in-memory DetectionDatasets is not supported."
+            )
+
+        images_in_memory = {}
+        for dataset in dataset_list:
+            images_in_memory.update(dataset._images_in_memory)
+
+        image_paths = list(
+            chain.from_iterable(dataset.image_paths for dataset in dataset_list)
+        )
+        image_paths_unique = list(dict.fromkeys(image_paths))
+        if len(image_paths) != len(image_paths_unique):
+            duplicates = find_duplicates(image_paths)
+            raise ValueError(
+                f"Image paths {duplicates} are not unique across datasets."
+            )
+        image_paths = image_paths_unique
+
+        classes = merge_class_lists(
+            class_lists=[dataset.classes for dataset in dataset_list]
+        )
+
+        annotations = {}
+        for dataset in dataset_list:
+            annotations.update(dataset.annotations)
+        for dataset in dataset_list:
+            class_index_mapping = build_class_index_mapping(
+                source_classes=dataset.classes, target_classes=classes
+            )
+            for image_path in dataset.image_paths:
+                annotations[image_path] = map_detections_class_id(
+                    source_to_target_mapping=class_index_mapping,
+                    detections=annotations[image_path],
+                )
+
+        merged_dataset = cls(
+            classes=classes,
+            images=image_paths,
+            annotations=annotations,
+        )
+        if all_in_memory:
+            merged_dataset._images_in_memory = images_in_memory
+        return merged_dataset
+
+    def as_pascal_voc(
+        self,
+        images_directory_path: str | None = None,
+        annotations_directory_path: str | None = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+        show_progress: bool = False,
+    ) -> None:
+        """
+        Exports the dataset to PASCAL VOC format. This method saves the images
+        and their corresponding annotations in PASCAL VOC format. Both output
+        layouts are preflighted before any files are written so a collision in
+        either target fails without partial output.
+
+        Args:
+            images_directory_path: The path to the directory
+                where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path: The path to
+                the directory where the annotations in PASCAL VOC format should be
+                saved. If not provided, annotations will not be saved.
+            min_image_area_percentage: The minimum percentage of
+                detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            max_image_area_percentage: The maximum percentage
+                of detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            approximation_percentage: The percentage of
+                polygon points to be removed from the input polygon,
+                in the range [0, 1). Argument is used only for segmentation datasets.
+            show_progress: If True, display a progress bar during saving.
+
+        Raises:
+            ValueError: If two image paths share the same basename (when
+                images_directory_path is set) or the same stem (when
+                annotations_directory_path is set), which would cause one
+                output file to overwrite another. Rename images to ensure
+                unique basenames before exporting a merged dataset.
+        """
+        if images_directory_path:
+            check_no_basename_collisions(
+                image_paths=self.image_paths,
+                key=lambda image_path: Path(image_path).name,
+                output_kind="image",
+            )
+        if annotations_directory_path:
+            check_no_basename_collisions(
+                image_paths=self.image_paths,
+                key=lambda image_path: f"{Path(image_path).stem}.xml",
+                output_kind="Pascal VOC annotation",
+            )
+
+        if images_directory_path:
+            save_dataset_images(
+                dataset=self,
+                images_directory_path=images_directory_path,
+                show_progress=show_progress,
+            )
+        if annotations_directory_path:
+            save_pascal_voc_annotations(
+                dataset=self,
+                annotations_directory_path=annotations_directory_path,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+                show_progress=show_progress,
+            )
+
+    @classmethod
+    def from_pascal_voc(
+        cls,
+        images_directory_path: str,
+        annotations_directory_path: str,
+        force_masks: bool = False,
+        show_progress: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from PASCAL VOC formatted data.
+
+        Args:
+            images_directory_path: Path to the directory containing the images.
+            annotations_directory_path: Path to the directory
+                containing the PASCAL VOC XML annotations.
+            force_masks: If True, forces masks to
+                be loaded for all annotations, regardless of whether they are present.
+            show_progress: If True, display a progress bar during loading.
+
+        Returns:
+            A DetectionDataset instance containing
+                the loaded images and annotations.
+
+        Examples:
+            ```python
+            import roboflow
+            from roboflow import Roboflow
+            import supervision as sv
+
+            roboflow.login()
+
+            rf = Roboflow()
+
+            project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
+            dataset = project.version(PROJECT_VERSION).download("voc")
+
+            ds = sv.DetectionDataset.from_pascal_voc(
+                images_directory_path=f"{dataset.location}/train/images",
+                annotations_directory_path=f"{dataset.location}/train/labels",
+                # pass show_progress=True to enable a tqdm progress bar
+            )
+
+            ds.classes
+            # ['dog', 'person']
+            ```
+        """
+
+        classes, image_paths, annotations = load_pascal_voc_annotations(
+            images_directory_path=images_directory_path,
+            annotations_directory_path=annotations_directory_path,
+            force_masks=force_masks,
+            show_progress=show_progress,
+        )
+
+        return DetectionDataset(
+            classes=classes, images=image_paths, annotations=annotations
+        )
+
+    @classmethod
+    def from_yolo(
+        cls,
+        images_directory_path: str,
+        annotations_directory_path: str,
+        data_yaml_path: str,
+        force_masks: bool = False,
+        is_obb: bool = False,
+        show_progress: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from YOLO formatted data.
+
+        Args:
+            images_directory_path: The path to the
+                directory containing the images.
+            annotations_directory_path: The path to the directory
+                containing the YOLO annotation files.
+            data_yaml_path: The path to the data
+                YAML file containing class information.
+            force_masks: If True, forces
+                masks to be loaded for all annotations,
+                regardless of whether they are present.
+            is_obb: If True, loads the annotations in OBB format.
+                OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
+                where pairs of [x, y] are box corners.
+            show_progress: If True, display a progress bar during loading.
+
+        Returns:
+            A DetectionDataset instance
+                containing the loaded images and annotations.
+
+        Examples:
+            ```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",
+                # pass show_progress=True to enable a tqdm progress bar
+            )
+
+            ds.classes
+            # ['dog', 'person']
+            ```
+        """
+        classes, image_paths, annotations = load_yolo_annotations(
+            images_directory_path=images_directory_path,
+            annotations_directory_path=annotations_directory_path,
+            data_yaml_path=data_yaml_path,
+            force_masks=force_masks,
+            is_obb=is_obb,
+            show_progress=show_progress,
+        )
+        return DetectionDataset(
+            classes=classes, images=image_paths, annotations=annotations
+        )
+
+    def as_yolo(
+        self,
+        images_directory_path: str | None = None,
+        annotations_directory_path: str | None = None,
+        data_yaml_path: str | None = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+        is_obb: bool = False,
+        show_progress: bool = False,
+    ) -> None:
+        """
+        Exports the dataset to YOLO format. This method saves the
+        images and their corresponding annotations in YOLO format.
+
+        Args:
+            images_directory_path: The path to the
+                directory where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path: 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: The path where the data.yaml
+                file should be saved.
+                If not provided, the file will not be saved.
+            min_image_area_percentage: The minimum percentage of
+                detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            max_image_area_percentage: The maximum percentage
+                of detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            approximation_percentage: The percentage of polygon points to
+                be removed from the input polygon, in the range [0, 1).
+                This is useful for simplifying the annotations.
+                Argument is used only for segmentation datasets.
+            is_obb: If True, exports annotations in OBB format
+                (`class_id x1 y1 x2 y2 x3 y3 x4 y4`) using the oriented
+                corners stored in `detections.data["xyxyxyxy"]`. Mirrors
+                `from_yolo(..., is_obb=True)`. Masks are ignored when
+                `is_obb=True`.
+            show_progress: If True, display a progress bar during saving.
+
+        Raises:
+            ValueError: If two image paths share the same basename (when
+                images_directory_path is set) or the same annotation
+                file name (when annotations_directory_path is set),
+                which would cause one output file to overwrite another.
+        """
+        if is_obb and (
+            min_image_area_percentage != 0.0
+            or max_image_area_percentage != 1.0
+            or approximation_percentage != 0.0
+        ):
+            import warnings
+
+            warnings.warn(
+                "`min_image_area_percentage`, `max_image_area_percentage`, and "
+                "`approximation_percentage` have no effect when `is_obb=True`; "
+                "OBB annotations use corner coordinates directly.",
+                UserWarning,
+                stacklevel=2,
+            )
+        # Pre-flight: validate output uniqueness before writing any file
+        if images_directory_path:
+            check_no_basename_collisions(
+                image_paths=self.image_paths,
+                key=lambda image_path: Path(image_path).name,
+                output_kind="image",
+            )
+        if annotations_directory_path:
+            check_no_basename_collisions(
+                image_paths=self.image_paths,
+                key=lambda image_path: Path(image_path).stem + ".txt",
+                output_kind="YOLO annotation",
+            )
+
+        if images_directory_path is not None:
+            save_dataset_images(
+                dataset=self,
+                images_directory_path=images_directory_path,
+                show_progress=show_progress,
+            )
+        if annotations_directory_path is not None:
+            save_yolo_annotations(
+                dataset=self,
+                annotations_directory_path=annotations_directory_path,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+                show_progress=show_progress,
+                is_obb=is_obb,
+            )
+        if data_yaml_path is not None:
+            save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)
+
+    @classmethod
+    def from_labelme(
+        cls,
+        images_directory_path: str,
+        annotations_directory_path: str,
+        force_masks: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from LabelMe formatted data.
+
+        LabelMe stores one JSON file per image, each containing a list of
+        ``shapes``. ``rectangle`` shapes are loaded as bounding boxes and
+        ``polygon`` shapes as masks (with their bounding boxes); other shape
+        types are skipped. Class names are inferred from the labels present in
+        the files. When an image file contains a ``polygon`` shape, or when
+        ``force_masks=True`` is set, both ``rectangle`` and ``polygon`` shapes
+        produce masks: rectangles via a four-corner polygon fill.
+
+        Args:
+            images_directory_path: The path to the
+                directory containing the images.
+            annotations_directory_path: The path to the directory
+                containing the LabelMe ``.json`` annotation files.
+            force_masks: If True, forces masks to be loaded for all
+                annotations, regardless of whether polygon shapes are present.
+                Requires ``imageWidth`` and ``imageHeight`` in every JSON file.
+
+        Returns:
+            A DetectionDataset instance containing
+                the loaded images and annotations.
+
+        Raises:
+            ValueError: If an annotation is malformed - for example
+                ``imagePath`` is empty or resolves to ``..``, a shape is
+                missing its ``label`` or ``points``, or a mask is required but
+                ``imageWidth`` / ``imageHeight`` are missing or zero.
+
+        Examples:
+            ```python
+            import supervision as sv
+
+            ds = sv.DetectionDataset.from_labelme(
+                images_directory_path="<IMAGES_DIRECTORY_PATH>",
+                annotations_directory_path="<ANNOTATIONS_DIRECTORY_PATH>",
+            )
+
+            ds.classes
+            # ['dog', 'person']
+            ```
+        """
+        classes, image_paths, annotations = load_labelme_annotations(
+            images_directory_path=images_directory_path,
+            annotations_directory_path=annotations_directory_path,
+            force_masks=force_masks,
+        )
+        return DetectionDataset(
+            classes=classes, images=image_paths, annotations=annotations
+        )
+
+    def as_labelme(
+        self,
+        images_directory_path: str | None = None,
+        annotations_directory_path: str | None = None,
+    ) -> None:
+        """
+        Exports the dataset to LabelMe format. This method saves the images and
+        their corresponding annotations as per-image LabelMe ``.json`` files.
+        Masked detections are written as ``polygon`` shapes whose vertices
+        approximate the mask contour, so masks are not bit-exact on round-trip.
+        Because the bounding box is recomputed from the quantized polygon contour
+        on re-import, bounding boxes for masked detections may also shift by
+        approximately one pixel after a save-load cycle.
+
+        Args:
+            images_directory_path: The path to the directory
+                where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_directory_path: The path to the directory where the
+                LabelMe ``.json`` files should be saved.
+                If not provided, annotations will not be saved.
+
+        Examples:
+            ```python
+            import supervision as sv
+
+            ds = sv.DetectionDataset(...)
+
+            ds.as_labelme(
+                images_directory_path="<IMAGES_DIRECTORY_PATH>",
+                annotations_directory_path="<ANNOTATIONS_DIRECTORY_PATH>",
+            )
+            ```
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                dataset=self, images_directory_path=images_directory_path
+            )
+        if annotations_directory_path is not None:
+            save_labelme_annotations(
+                dataset=self,
+                annotations_directory_path=annotations_directory_path,
+            )
+
+    @classmethod
+    def from_createml(
+        cls,
+        images_directory_path: str,
+        annotations_path: str,
+        show_progress: bool = False,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from CreateML formatted data.
+
+        CreateML stores object-detection annotations in a single JSON file as a
+        list of per-image entries, with each box expressed as a pixel-space
+        centre point plus width and height. Class names are inferred from the
+        labels present in the file.
+
+        Args:
+            images_directory_path: The path to the directory containing the
+                images.
+            annotations_path: The path to the CreateML json annotation file.
+            show_progress: If True, display a progress bar during loading.
+
+        Returns:
+            A DetectionDataset instance containing the loaded images and
+            annotations.
+
+        Examples:
+            ```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("createml")
+
+            ds = sv.DetectionDataset.from_createml(
+                images_directory_path=f"{dataset.location}/train",
+                annotations_path=f"{dataset.location}/train/_annotations.createml.json",
+            )
+
+            ds.classes
+            # ['dog', 'person']
+            ```
+        """
+        classes, image_paths, annotations = load_createml_annotations(
+            images_directory_path=images_directory_path,
+            annotations_path=annotations_path,
+            show_progress=show_progress,
+        )
+        return DetectionDataset(
+            classes=classes, images=image_paths, annotations=annotations
+        )
+
+    def as_createml(
+        self,
+        images_directory_path: str | None = None,
+        annotations_path: str | None = None,
+        show_progress: bool = False,
+    ) -> None:
+        """
+        Exports the dataset to CreateML format. This method saves the
+        images and their corresponding annotations in CreateML format.
+
+        Args:
+            images_directory_path: The path to the directory where the images
+                should be saved. If not provided, images will not be saved.
+            annotations_path: The path to the CreateML json annotation file.
+                If not provided, the annotations will not be saved.
+            show_progress: If True, display a progress bar while saving images.
+
+        Returns:
+            None. Side-effects only: writes images and/or annotation file.
+
+        Examples:
+            ```python
+            import supervision as sv
+
+            ds = sv.DetectionDataset(classes=["dog"], images=[], annotations={})
+            ds.as_createml(
+                images_directory_path="/tmp/images",
+                annotations_path="/tmp/annotations.json",
+            )
+            ```
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                dataset=self,
+                images_directory_path=images_directory_path,
+                show_progress=show_progress,
+            )
+        if annotations_path is not None:
+            save_createml_annotations(
+                dataset=self,
+                annotations_path=annotations_path,
+            )
+
+    @classmethod
+    def from_coco(
+        cls,
+        images_directory_path: str,
+        annotations_path: str,
+        force_masks: bool = False,
+        show_progress: bool = False,
+        *,
+        use_iscrowd: bool = True,
+    ) -> DetectionDataset:
+        """
+        Creates a Dataset instance from COCO formatted data.
+
+        Args:
+            images_directory_path: The path to the
+                directory containing the images.
+            annotations_path: The path to the json annotation files.
+            force_masks: If True,
+                forces masks to be loaded for all annotations,
+                regardless of whether they are present.
+            show_progress: If True, display a progress bar during loading.
+            use_iscrowd: If True, includes COCO ``iscrowd`` and ``area``
+                annotation fields in ``Detections.data``.
+        Returns:
+            A DetectionDataset instance containing
+                the loaded images and annotations.
+
+        Examples:
+            ```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("coco")
+
+            ds = sv.DetectionDataset.from_coco(
+                images_directory_path=f"{dataset.location}/train",
+                annotations_path=f"{dataset.location}/train/_annotations.coco.json",
+                # pass show_progress=True to enable a tqdm progress bar
+            )
+
+            ds.classes
+            # ['dog', 'person']
+            ```
+        """
+        classes, images, annotations = load_coco_annotations(
+            images_directory_path=images_directory_path,
+            annotations_path=annotations_path,
+            force_masks=force_masks,
+            use_iscrowd=use_iscrowd,
+            show_progress=show_progress,
+        )
+        return DetectionDataset(classes=classes, images=images, annotations=annotations)
+
+    def as_coco(
+        self,
+        images_directory_path: str | None = None,
+        annotations_path: str | None = None,
+        min_image_area_percentage: float = 0.0,
+        max_image_area_percentage: float = 1.0,
+        approximation_percentage: float = 0.0,
+        starting_image_id: int = 1,
+        starting_annotation_id: int = 1,
+        show_progress: bool = False,
+    ) -> tuple[int, int]:
+        """
+        Exports the dataset to COCO format. This method saves the
+        images and their corresponding annotations in COCO format.
+
+        !!! tip
+
+            The format of the mask is determined automatically based on its structure:
+
+            - If a mask contains multiple disconnected components or holes, it will be
+            saved using the Run-Length Encoding (RLE) format for efficient storage and
+            processing.
+            - If a mask consists of a single, contiguous region without any holes, it
+            will be encoded as a polygon, preserving the outline of the object.
+
+            This automatic selection ensures that the masks are stored in the most
+            appropriate and space-efficient format, complying with COCO dataset
+            standards.
+
+        Args:
+            images_directory_path: The path to the directory
+                where the images should be saved.
+                If not provided, images will not be saved.
+            annotations_path: The path to COCO annotation file.
+            min_image_area_percentage: The minimum percentage of
+                detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            max_image_area_percentage: The maximum percentage of
+                detection area relative to
+                the image area for a detection to be included.
+                Argument is used only for segmentation datasets.
+            approximation_percentage: The percentage of polygon points
+                to be removed from the input polygon,
+                in the range [0, 1). This is useful for simplifying the annotations.
+                Argument is used only for segmentation datasets.
+            starting_image_id: First image id to assign in the exported file.
+                Defaults to ``1``. Override when exporting multiple splits into
+                a coordinated COCO collection so ids remain unique across the
+                set (see example below).
+            starting_annotation_id: First annotation id to assign in the
+                exported file. Defaults to ``1``. Override for the same
+                multi-split reason as ``starting_image_id``.
+            show_progress: If True, display a progress bar during saving.
+
+        Returns:
+            A ``(next_image_id, next_annotation_id)`` tuple containing the
+            first unused ids after this export. Feed them straight back into
+            ``starting_image_id`` and ``starting_annotation_id`` on the next
+            split so ids stay globally unique. When ``annotations_path`` is
+            ``None`` (images-only export) the starting ids are returned
+            unchanged so chaining still composes.
+
+        Example:
+            ```python
+            # Exporting train, valid, and test splits with non-colliding ids
+            # so the three annotation files can later be merged into one COCO.
+            next_image_id, next_annotation_id = train_ds.as_coco(
+                images_directory_path="out/train/images",
+                annotations_path="out/train/annotations.json",
+            )
+            next_image_id, next_annotation_id = valid_ds.as_coco(
+                images_directory_path="out/valid/images",
+                annotations_path="out/valid/annotations.json",
+                starting_image_id=next_image_id,
+                starting_annotation_id=next_annotation_id,
+            )
+            _, _ = test_ds.as_coco(
+                images_directory_path="out/test/images",
+                annotations_path="out/test/annotations.json",
+                starting_image_id=next_image_id,
+                starting_annotation_id=next_annotation_id,
+            )  # return value not needed — no further split
+            ```
+        """
+        if images_directory_path is not None:
+            save_dataset_images(
+                dataset=self,
+                images_directory_path=images_directory_path,
+                show_progress=show_progress,
+            )
+        if annotations_path is not None:
+            return save_coco_annotations(
+                dataset=self,
+                annotation_path=annotations_path,
+                min_image_area_percentage=min_image_area_percentage,
+                max_image_area_percentage=max_image_area_percentage,
+                approximation_percentage=approximation_percentage,
+                starting_image_id=starting_image_id,
+                starting_annotation_id=starting_annotation_id,
+                show_progress=show_progress,
+            )
+        return starting_image_id, starting_annotation_id
+
+
+@dataclass
+class ClassificationDataset(BaseDataset):
+    """
+    Contains information about a classification dataset, handles lazy image
+    loading, dataset splitting.
+
+    Attributes:
+        classes: List containing dataset class names.
+        images:
+            List of image paths or dictionary mapping image name to image data.
+        annotations: Dictionary mapping
+            image name to annotations.
+    """
+
+    def __init__(
+        self,
+        classes: list[str],
+        images: list[str] | dict[str, npt.NDArray[np.uint8]],
+        annotations: dict[str, Classifications],
+    ) -> None:
+        self.classes = classes
+
+        if set(images) != set(annotations):
+            raise ValueError(
+                "The keys of the images and annotations dictionaries must match."
+            )
+        self.annotations = annotations
+
+        # Eliminate duplicates while preserving order
+        self.image_paths = list(dict.fromkeys(images))
+
+        self._images_in_memory: dict[str, npt.NDArray[np.uint8]] = {}
+        if isinstance(images, dict):
+            self._images_in_memory = images
+            warn_deprecated(
+                "Passing a `Dict[str, np.ndarray]` into `ClassificationDataset` is "
+                "deprecated and will be removed in a future release. Use "
+                "a list of paths `List[str]` instead."
+            )
+
+    def _get_image(self, image_path: str) -> npt.NDArray[np.uint8]:
+        """Assumes that image is in dataset."""
+        if self._images_in_memory:
+            return self._images_in_memory[image_path]
+        image = cv2.imread(image_path)
+        if image is None:
+            raise ValueError(f"Could not read image from path: {image_path}")
+        return cast(npt.NDArray[np.uint8], image)
+
+    def __len__(self) -> int:
+        return len(self._images_in_memory) or len(self.image_paths)
+
+    def __getitem__(self, i: int) -> tuple[str, npt.NDArray[np.uint8], Classifications]:
+        """
+        Returns:
+            The image path, image data,
+                and its corresponding annotation at index i.
+        """
+        image_path = self.image_paths[i]
+        image = self._get_image(image_path)
+        annotation = self.annotations[image_path]
+        return image_path, image, annotation
+
+    def __iter__(
+        self,
+    ) -> Iterator[tuple[str, npt.NDArray[np.uint8], Classifications]]:
+        """
+        Iterate over the images and annotations in the dataset.
+
+        Yields:
+            Tuples containing the image path, image data, and its annotation.
+        """
+        for i in range(len(self)):
+            image_path, image, annotation = self[i]
+            yield image_path, image, annotation
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, ClassificationDataset):
+            return False
+
+        if self.classes != other.classes:
+            return False
+
+        if self.image_paths != other.image_paths:
+            return False
+
+        if self._images_in_memory or other._images_in_memory:
+            if not np.array_equal(
+                list(self._images_in_memory.values()),
+                list(other._images_in_memory.values()),
+            ):
+                return False
+
+        if self.annotations != other.annotations:
+            return False
+
+        return True
+
+    def split(
+        self,
+        split_ratio: float = 0.8,
+        random_state: int | None = None,
+        shuffle: bool = True,
+    ) -> tuple[ClassificationDataset, ClassificationDataset]:
+        """
+        Splits the dataset into two parts (training and testing)
+            using the provided split_ratio.
+
+        Args:
+            split_ratio: The ratio of the training
+                set to the entire dataset.
+            random_state: The seed for the
+                random number generator. This is used for reproducibility.
+            shuffle: Whether to shuffle the data before splitting.
+
+        Returns:
+            A tuple containing
+            the training and testing datasets.
+
+        Examples:
+            ```pycon
+            >>> import numpy as np
+            >>> import supervision as sv
+            >>> cd = sv.ClassificationDataset(
+            ...     classes=['cat', 'dog'],
+            ...     images={
+            ...         'img1.jpg': np.zeros((100, 100, 3), dtype=np.uint8),
+            ...         'img2.jpg': np.zeros((100, 100, 3), dtype=np.uint8),
+            ...     },
+            ...     annotations={
+            ...         'img1.jpg': sv.Classifications(class_id=np.array([0])),
+            ...         'img2.jpg': sv.Classifications(class_id=np.array([1])),
+            ...     }
+            ... )
+            >>> train_cd, test_cd = cd.split(split_ratio=0.5, random_state=42)
+            >>> len(train_cd), len(test_cd)
+            (1, 1)
+
+            ```
+        """
+        train_paths, test_paths = train_test_split(
+            data=self.image_paths,
+            train_ratio=split_ratio,
+            random_state=random_state,
+            shuffle=shuffle,
+        )
+
+        train_input: list[str] | dict[str, npt.NDArray[np.uint8]]
+        test_input: list[str] | dict[str, npt.NDArray[np.uint8]]
+        if self._images_in_memory:
+            train_input = {path: self._images_in_memory[path] for path in train_paths}
+            test_input = {path: self._images_in_memory[path] for path in test_paths}
+        else:
+            train_input = train_paths
+            test_input = test_paths
+        train_annotations = {path: self.annotations[path] for path in train_paths}
+        test_annotations = {path: self.annotations[path] for path in test_paths}
+
+        train_dataset = ClassificationDataset(
+            classes=self.classes,
+            images=train_input,
+            annotations=train_annotations,
+        )
+        test_dataset = ClassificationDataset(
+            classes=self.classes,
+            images=test_input,
+            annotations=test_annotations,
+        )
+
+        return train_dataset, test_dataset
+
+    def as_folder_structure(
+        self, root_directory_path: str, show_progress: bool = False
+    ) -> None:
+        """
+        Saves the dataset as a multi-class folder structure.
+
+        Args:
+            root_directory_path: The path to the directory
+                where the dataset will be saved.
+            show_progress: If True, display a progress bar during saving.
+        """
+        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_save_path, image, annotation in tqdm(
+            self,
+            total=len(self),
+            desc="Saving classification images",
+            disable=not show_progress,
+        ):
+            image_name = Path(image_save_path).name
+            class_id = (
+                annotation.class_id[0]
+                if annotation.confidence is None
+                else annotation.get_top_k(1)[0][0]
+            )
+            class_name = self.classes[class_id]
+            image_save_path = os.path.join(root_directory_path, class_name, image_name)
+            cv2.imwrite(image_save_path, image)
+
+    @classmethod
+    def from_folder_structure(
+        cls, root_directory_path: str, show_progress: bool = False
+    ) -> ClassificationDataset:
+        """
+        Load data from a multiclass folder structure into a ClassificationDataset.
+
+        Args:
+            root_directory_path: The path to the dataset directory. Hidden
+                entries, root-level files, nested directories, and files whose
+                suffix is not a supported image extension are ignored.
+            show_progress: If True, display a progress bar during loading.
+
+        Returns:
+            The dataset.
+
+        Examples:
+            ```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"
+                # pass show_progress=True to enable a tqdm progress bar
+            )
+            ```
+        """
+        root_directory = Path(root_directory_path)
+        classes = sorted(
+            {
+                entry.name
+                for entry in root_directory.iterdir()
+                if entry.is_dir() and not entry.name.startswith(".")
+            }
+        )
+
+        image_paths = []
+        annotations = {}
+
+        for class_id, class_name in enumerate(
+            tqdm(
+                classes,
+                total=len(classes),
+                desc="Loading classification dataset",
+                disable=not show_progress,
+            )
+        ):
+            class_directory = root_directory / class_name
+
+            for image_path in sorted(
+                class_directory.iterdir(), key=lambda path: path.name
+            ):
+                if image_path.name.startswith(".") or not image_path.is_file():
+                    continue
+                if image_path.suffix.lower() not in _IMAGE_FILE_EXTENSIONS:
+                    continue
+                image_paths.append(str(image_path))
+                annotations[str(image_path)] = Classifications(
+                    class_id=np.array([class_id]),
+                )
+
+        return cls(
+            classes=classes,
+            images=image_paths,
+            annotations=annotations,
+        )