diff --git a/supervision/dataset/core.py b/supervision/dataset/core.py index 7d320bbc9..4ff6fab0c 100644 --- a/supervision/dataset/core.py +++ b/supervision/dataset/core.py @@ -504,6 +504,7 @@ def as_yolo( images_directory_path: Optional[str] = None, annotations_directory_path: Optional[str] = None, data_yaml_path: Optional[str] = None, + subset_type: Optional[str] = None, min_image_area_percentage: float = 0.0, max_image_area_percentage: float = 1.0, approximation_percentage: float = 0.0, @@ -523,6 +524,11 @@ def as_yolo( data_yaml_path (Optional[str]): The path where the data.yaml file should be saved. If not provided, the file will not be saved. + subset_type (Optional[str]): The subset type of the exported + dataset. + Accepted values: "train", "val", "test" + If not provided, the subset dataset path will not be + added or modified in data yaml file. min_image_area_percentage (float): The minimum percentage of detection area relative to the image area for a detection to be included. @@ -549,7 +555,11 @@ def as_yolo( approximation_percentage=approximation_percentage, ) if data_yaml_path is not None: - save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes) + save_data_yaml( + data_yaml_path=data_yaml_path, + classes=self.classes, + subset_type=subset_type, + ) @classmethod def from_coco( diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 0ecbac4b5..fdd497262 100644 --- a/supervision/dataset/formats/yolo.py +++ b/supervision/dataset/formats/yolo.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union import cv2 import numpy as np @@ -266,7 +266,16 @@ def save_yolo_annotations( 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} +def save_data_yaml( + data_yaml_path: str, classes: List[str], subset_type: Union[str, None] +) -> None: + if Path(data_yaml_path).exists(): + data = read_yaml_file(data_yaml_path) + else: + data = {} + data["nc"] = len(classes) + data["names"] = classes + if subset_type is not None: + data[subset_type] = f"{subset_type}/images" Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True) save_yaml_file(data=data, file_path=data_yaml_path) diff --git a/supervision/utils/file.py b/supervision/utils/file.py index 987bf64f3..4627cc614 100644 --- a/supervision/utils/file.py +++ b/supervision/utils/file.py @@ -112,7 +112,6 @@ def save_json_file(data: dict, file_path: Union[str, Path], indent: int = 3) -> Write a dict to a json file. Args: - indent: data (dict): dict with unique keys and value as pair. file_path (Union[str, Path]): The file path as a string or Path object. """ @@ -137,10 +136,9 @@ def read_yaml_file(file_path: Union[str, Path]) -> Dict: def save_yaml_file(data: dict, file_path: Union[str, Path]) -> None: """ - Save a dict to a json file. + Save a dict to a yaml file. Args: - indent: data (dict): dict with unique keys and value as pair. file_path (Union[str, Path]): The file path as a string or Path object. """