Add train, val, test folder paths in data.yaml at save_data_yaml()

supervision/dataset/core.py

@@ -510,6 +510,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,
@@ -529,6 +530,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.
@@ -555,7 +561,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(

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 numpy as np
 from PIL import Image
@@ -283,7 +283,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)

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.
     """