from __future__ import annotations __all__ = ["check_no_basename_collisions", "train_test_split"] import copy import os import random import shutil from collections.abc import Callable from pathlib import Path from typing import TYPE_CHECKING, TypeVar, cast import numpy as np import numpy.typing as npt from deprecate import deprecated, void # type: ignore[import-untyped,unused-ignore] from tqdm.auto import tqdm from supervision import _cv2 as cv2 from supervision.detection.core import Detections from supervision.detection.utils.converters import mask_to_polygons from supervision.detection.utils.converters import ( mask_to_rle as _mask_to_rle, ) from supervision.detection.utils.converters import ( rle_to_mask as _rle_to_mask, ) from supervision.detection.utils.polygons import ( approximate_polygon, filter_polygons_by_area, ) @deprecated(target=_mask_to_rle, deprecated_in="0.28.0", remove_in="0.31.0") # type: ignore[untyped-decorator] def mask_to_rle( mask: npt.NDArray[np.bool_], compressed: bool = False ) -> list[int] | str: """Deprecated since 0.28.0. Use `supervision.detection.utils.converters.mask_to_rle`. """ return cast(list[int] | str, void(mask, compressed)) @deprecated(target=_rle_to_mask, deprecated_in="0.28.0", remove_in="0.31.0") # type: ignore[untyped-decorator] def rle_to_mask( rle: npt.NDArray[np.integer] | list[int] | str | bytes, resolution_wh: tuple[int, int], ) -> npt.NDArray[np.bool_]: """Deprecated since 0.28.0. Use `supervision.detection.utils.converters.rle_to_mask`. """ return cast(npt.NDArray[np.bool_], void(rle, resolution_wh)) if TYPE_CHECKING: from supervision.dataset.core import DetectionDataset T = TypeVar("T") def approximate_mask_with_polygons( mask: npt.NDArray[np.bool_], min_image_area_percentage: float = 0.0, max_image_area_percentage: float = 1.0, approximation_percentage: float = 0.0, ) -> list[npt.NDArray[np.number]]: """Filter mask polygons by area and optionally simplify them. The default `approximation_percentage=0.0` preserves the original contour unless callers explicitly ask for simplification. """ height, width = mask.shape image_area = height * width minimum_detection_area = min_image_area_percentage * image_area maximum_detection_area = max_image_area_percentage * image_area polygons = cast(list[npt.NDArray[np.number]], mask_to_polygons(mask=mask)) if len(polygons) == 1: polygons = filter_polygons_by_area( polygons=polygons, min_area=None, max_area=maximum_detection_area ) else: polygons = filter_polygons_by_area( polygons=polygons, min_area=minimum_detection_area, max_area=maximum_detection_area, ) return [ approximate_polygon(polygon=polygon, percentage=approximation_percentage) for polygon in polygons ] def merge_class_lists(class_lists: list[list[str]]) -> list[str]: unique_classes = set() for class_list in class_lists: for class_name in class_list: unique_classes.add(class_name) return sorted(list(unique_classes)) def build_class_index_mapping( source_classes: list[str], target_classes: list[str] ) -> dict[int, int]: """Returns the index map of source classes -> target classes.""" index_mapping = {} for i, class_name in enumerate(source_classes): if class_name not in target_classes: raise ValueError( f"Class {class_name} not found in target classes. " "source_classes must be a subset of target_classes." ) corresponding_index = target_classes.index(class_name) index_mapping[i] = corresponding_index return index_mapping def map_detections_class_id( source_to_target_mapping: dict[int, int], detections: Detections ) -> Detections: if detections.class_id is None: raise ValueError("Detections must have class_id attribute.") if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()): raise ValueError( "Detections class_id must be a subset of source_to_target_mapping keys." ) detections_copy = copy.deepcopy(detections) if len(detections) > 0: detections_copy.class_id = np.vectorize(source_to_target_mapping.get)( detections_copy.class_id ) return detections_copy def check_no_basename_collisions( image_paths: list[str], key: Callable[[str], str], output_kind: str, ) -> None: """Raise if two image paths would be written to the same output file. Dataset image paths may share a basename when they originate from different directories (a legal, common state after :meth:`DetectionDataset.merge`). Exporting them into a single flat output directory keyed on the basename or stem would silently overwrite one file with another and mispair images with their annotations. This guard detects such collisions before any file is written and names the colliding source paths. Args: image_paths: The dataset image paths about to be written. key: Maps an image path to the output file name it would be written to. output_kind: Human-readable description of the output (e.g. ``"image"`` or ``"YOLO annotation"``) used in the error message. Raises: ValueError: If two image paths map to the same output file name. Examples: ```pycon >>> from pathlib import Path >>> from supervision.dataset.utils import check_no_basename_collisions >>> check_no_basename_collisions( ... ["a/img.jpg", "b/img.jpg"], lambda p: Path(p).name, "image" ... ) Traceback (most recent call last): ... ValueError: Cannot export dataset: image paths 'a/img.jpg' and ... ``` """ seen: dict[str, tuple[str, str]] = {} # casefold(key) → (original name, image_path) for image_path in image_paths: output_name = key(image_path) case_key = output_name.casefold() if case_key in seen: first_name, first_path = seen[case_key] raise ValueError( f"Cannot export dataset: image paths {first_path!r} and " f"{image_path!r} both map to {output_kind} file {first_name!r}. " "Ensure all image basenames are unique before exporting." ) seen[case_key] = (output_name, image_path) def save_dataset_images( dataset: DetectionDataset, images_directory_path: str, show_progress: bool = False, ) -> None: """Save all images from a dataset to a directory. Images already in memory are written with ``cv2.imwrite``; images stored only as file paths are copied with ``shutil.copyfile``. Args: dataset: The dataset whose images are saved. images_directory_path: Destination directory path; created automatically if it does not exist. show_progress: If ``True``, display a tqdm progress bar while saving images. Examples: ```pycon >>> from supervision.dataset.core import DetectionDataset >>> from supervision.dataset.utils import save_dataset_images >>> dataset = DetectionDataset(classes=["cat"], images={}, annotations={}) >>> save_dataset_images(dataset, "/tmp/images") ``` """ check_no_basename_collisions( image_paths=dataset.image_paths, key=lambda image_path: Path(image_path).name, output_kind="image", ) Path(images_directory_path).mkdir(parents=True, exist_ok=True) for image_path in tqdm( dataset.image_paths, desc="Saving images", disable=not show_progress, ): final_path = os.path.join(images_directory_path, Path(image_path).name) if image_path in dataset._images_in_memory: image = dataset._images_in_memory[image_path] cv2.imwrite(final_path, image) else: shutil.copyfile(image_path, final_path) def train_test_split( data: list[T], train_ratio: float = 0.8, random_state: int | None = None, shuffle: bool = True, ) -> tuple[list[T], list[T]]: """ Splits the data into two parts using the provided train_ratio. Args: data: The data to split. train_ratio: The ratio of the training set to the entire dataset. random_state: The seed for the random number generator. shuffle: Whether to shuffle the data before splitting. Returns: The split data. The input list is copied and never mutated. Examples: ```pycon >>> train, test = train_test_split( ... [1, 2, 3, 4, 5], train_ratio=0.6, random_state=0 ... ) >>> len(train), len(test) (3, 2) ``` """ rng = random.Random(random_state) # noqa: S311 — dataset split, not cryptographic if shuffle: data = list(data) rng.shuffle(data) else: data = list(data) # copy to guarantee non-mutation split_index = int(len(data) * train_ratio) return data[:split_index], data[split_index:]