diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 0682e7eb..7d320bbc 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -28,6 +28,7 @@ from supervision.dataset.utils import ( build_class_index_mapping, map_detections_class_id, merge_class_lists, + save_dataset_images, train_test_split, ) from supervision.detection.core import Detections @@ -356,7 +357,8 @@ class DetectionDataset(BaseDataset): in the range [0, 1). Argument is used only for segmentation datasets. """ if images_directory_path: - self._save_images( + save_dataset_images( + dataset=self, images_directory_path=images_directory_path, ) if annotations_directory_path: @@ -535,7 +537,9 @@ class DetectionDataset(BaseDataset): Argument is used only for segmentation datasets. """ if images_directory_path is not None: - self._save_images(images_directory_path=images_directory_path) + save_dataset_images( + dataset=self, images_directory_path=images_directory_path + ) if annotations_directory_path is not None: save_yolo_annotations( dataset=self, @@ -642,7 +646,9 @@ class DetectionDataset(BaseDataset): Argument is used only for segmentation datasets. """ if images_directory_path is not None: - self._save_images(images_directory_path=images_directory_path) + save_dataset_images( + dataset=self, images_directory_path=images_directory_path + ) if annotations_path is not None: save_coco_annotations( dataset=self, @@ -652,12 +658,6 @@ class DetectionDataset(BaseDataset): approximation_percentage=approximation_percentage, ) - def _save_images(self, images_directory_path: str) -> None: - Path(images_directory_path).mkdir(parents=True, exist_ok=True) - for image_path, image, _ in self: - final_path = os.path.join(images_directory_path, image_path) - cv2.imwrite(final_path, image) - @dataclass class ClassificationDataset(BaseDataset): diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py index ef59050d..4da19b62 100644 --- a/supervision/dataset/utils.py +++ b/supervision/dataset/utils.py @@ -1,7 +1,10 @@ import copy +import os import random -from typing import Dict, List, Optional, Tuple, TypeVar, Union +from pathlib import Path +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union +import cv2 import numpy as np import numpy.typing as npt @@ -12,6 +15,9 @@ from supervision.detection.utils import ( mask_to_polygons, ) +if TYPE_CHECKING: + from supervision.dataset.core import DetectionDataset + T = TypeVar("T") @@ -91,6 +97,15 @@ def map_detections_class_id( return detections_copy +def save_dataset_images( + dataset: "DetectionDataset", images_directory_path: str +) -> None: + Path(images_directory_path).mkdir(parents=True, exist_ok=True) + for image_path, image, _ in dataset: + final_path = os.path.join(images_directory_path, image_path) + cv2.imwrite(final_path, image) + + def train_test_split( data: List[T], train_ratio: float = 0.8,