YoloDatasetCustomizer 0.1.0__tar.gz → 0.3.0__tar.gz

yolodatasetcustomizer-0.3.0/.github/workflows/ci.yml

@@ -0,0 +1,35 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        # "3" resolves to the latest released stable CPython 3.x, so this
+        # tracks new releases automatically without needing manual bumps
+        python-version: ["3.10", "3"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --group dev
+
+      - name: Lint (ruff)
+        run: uv run ruff check src test
+
+      - name: Type-check (mypy)
+        run: uv run mypy
+
+      - name: Test (pytest)
+        run: uv run pytest -q

{yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: YoloDatasetCustomizer
-Version: 0.1.0
+Version: 0.3.0
 Summary: Load one or multiple datasets. Reduce, mix or merge into a custom dataset.
 Project-URL: Homepage, https://github.com/mLuca/YoloDatasetCustomizer
 Project-URL: Repository, https://github.com/mLuca/YoloDatasetCustomizer
@@ -58,6 +58,11 @@ The project uses uv. Install dependencies and run tests with:
 
     uv run pytest
 
+# Lint & type-check
+
+    uv run ruff check src test
+    uv run mypy
+
 # Build project
 
 Build the package

{yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.3.0}/README.md

@@ -31,6 +31,11 @@ The project uses uv. Install dependencies and run tests with:
 
     uv run pytest
 
+# Lint & type-check
+
+    uv run ruff check src test
+    uv run mypy
+
 # Build project
 
 Build the package

{yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "YoloDatasetCustomizer"
-version = "0.1.0"
+version = "0.3.0"
 description = "Load one or multiple datasets. Reduce, mix or merge into a custom dataset."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -40,4 +40,19 @@ Repository = "https://github.com/mLuca/YoloDatasetCustomizer"
 [dependency-groups]
 dev = [
     "pytest>=9.1.1",
+    "twine>=6.2.0",
+    "ruff>=0.14.0",
+    "mypy>=1.18.0",
 ]
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+ignore = ["E501"]
+
+[tool.mypy]
+python_version = "3.10"
+files = ["src"]
+ignore_missing_imports = true

yolodatasetcustomizer-0.3.0/resources/Improvements.txt

@@ -0,0 +1,57 @@
+Priority-ordered improvement list (last reviewed 2026-08-14)
+
+FIXED 1. Fix brittle path and OS handling
+   Why: The code constructs file paths with string concatenation and assumes Unix-style separators, which can break on Windows and in nested dataset layouts. Since this tool operates on filesystem paths, cross-platform reliability is a core requirement.
+
+FIXED 2. Remove mutable default arguments
+   Why: class_names: list[str] = [] creates shared state across instances. This is a Python bug pattern that can cause hard-to-trace behavior and incorrect dataset generation when objects are reused.
+
+FIXED 3. Fix broken internal state and debug output
+   Why: __repr__ references attributes that do not exist, which signals stale or inconsistent internal state. This is a maintainability problem and a sign that code may be harder to trust during debugging.
+
+FIXED 4. Refactor the monolithic implementation into smaller modules
+   Why: The main file handles reading YAML, writing YAML, copying labels, scanning directories, and building datasets all at once. Splitting responsibilities would reduce complexity, improve testability, and make future features safer to add.
+
+FIXED 1. Fix broken copy-failure error handling in customizer.py
+   Why: `if not shutil.copyfile(img, destination):` is dead code — copyfile returns a truthy destination path on success and raises on failure, so this branch never triggers and a real copy error (permissions, disk full, locked file) crashes the whole run with an unhandled exception instead of returning False with a clear message.
+
+FIXED 2. Fix the path-existence skip logic in create_new_dataset_for_class_names
+   Why: `for path in [old_labels_path, old_images_path]: if not os.path.exists(path): ...continue` only continues that 2-item inner loop, not the split-processing loop it's meant to guard. A missing labels/images directory is not actually skipped as intended, so processing silently proceeds against a non-existent path (glob just returns nothing, masking the real problem).
+
+FIXED 3. Standardize class-index typing and conversion
+   Why: Indices are handled as strings in some places and as numeric concepts in others. This inconsistency increases the chance of subtle label remapping bugs that are especially dangerous in object-detection datasets.
+
+4. Tighten validation and error handling generally
+   Why: add_data_sets() prints a warning and continues on bad YAML instead of surfacing a clear, actionable failure or strict mode. Silently skipping broken YAML or missing directories can produce partially valid outputs and make debugging much harder.
+
+5. Stop leaking mutable internal state from getters
+   Why: get_found_data_file_paths() and get_class_names() return the live internal set/list rather than a copy. Callers can accidentally mutate the object's internal state from outside, breaking encapsulation.
+
+FIXED 6. Add CI (GitHub Actions) for tests, lint, and type-checking
+   Why: The package classifies itself as "Production/Stable" and is published to PyPI, but nothing runs tests or checks automatically on push/PR. Bugs like items 1 and 2 above would have been easy to catch with basic CI.
+
+FIXED 7. Add lint/type-check tooling (ruff + mypy) to the dev dependency group
+   Why: There is currently no configured static analysis. Tooling like this would have flagged the always-true/dead `if not shutil.copyfile(...)` check and similar issues automatically.
+
+8. Replace print()-based warnings/errors with the `logging` module
+   Why: A library shouldn't write directly to stdout. Using `logging` lets consumers control verbosity, redirect output, and integrate the tool into larger pipelines.
+
+FIXED 9. Clean up naming and spelling inconsistencies
+   Why: The misspelling "indeces" is used in some names (copy_by_class_indeces, __generate_new_class_indeces) while the correct "indices" is used elsewhere (get_class_indices, get_indices_for_names) in the same codebase. This is not just cosmetic; it slows down maintenance and increases the chance of mistakes during refactors.
+
+10. Remove or clarify the directory-handling branch in LabelFile.__init__
+   Why: `_get_unique_path(file_path) if os.path.isdir(file_path) else ...` produces a non-existent path (by design of _get_unique_path) whenever a directory is passed, which then immediately raises FileNotFoundError. This looks like dead or misleading code with no clear purpose.
+
+11. Make SUPPORTED_IMAGE_FORMATS a class-level constant
+   Why: It's currently rebuilt as a public mutable list on every instance in __init__. A class-level tuple/frozenset avoids repeated allocation and accidental external mutation.
+
+12. Add a py.typed marker file
+   Why: The package ships inline type hints and is installable from PyPI, but without py.typed, downstream type checkers (mypy/pyright) will ignore those hints for consumers of the published package.
+
+13. Improve test coverage around edge cases
+   Why: The current tests cover the happy path well, plus a few negative cases, but not the bugs found above (missing labels/images dir, copy failures) or other realistic failures such as duplicate file names, empty split folders, and unexpected dataset layouts.
+
+14. Clarify and formally validate documented limitations
+   Why: The README already explains assumptions like fixed dataset naming and ignored YAML fields (`path`, `script`). Those constraints should be treated as explicit, validated API requirements (and paired with clearer usage examples) so users understand exactly what is supported and what is not.
+
+Highest priority overall: items 1 and 2 are real correctness bugs in the core dataset-generation path that can silently produce wrong, incomplete, or crashing output today. Fix those first, then invest in CI and lint tooling (items 6-7) so regressions like this are caught automatically going forward.

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/__init__.py

@@ -0,0 +1,17 @@
+"""YOLO Dataset Customizer - Load, mix, merge, and customize YOLO datasets."""
+
+from .customizer import YoloDatasetCustomizer
+from .labels import LabelFile
+from .reader import YoloDataFileReader
+from .writer import YoloDataFileWriter
+
+__version__ = "0.3.0"
+__author__ = "Luca Mazzon"
+__email__ = "mazzon.luca@gmail.com"
+
+__all__ = [
+    "YoloDatasetCustomizer",
+    "YoloDataFileWriter",
+    "YoloDataFileReader",
+    "LabelFile",
+]

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/customizer.py

@@ -0,0 +1,174 @@
+import glob
+import os
+import shutil
+
+import yaml
+
+from .labels import LabelFile
+from .paths import _SPLIT_NAMES, _get_unique_path, _normalize_path
+from .reader import YoloDataFileReader
+from .writer import YoloDataFileWriter
+
+
+class YoloDatasetCustomizer:
+    """Find and customize YOLO datasets across one or more directory paths."""
+
+    def __init__(self, data_set_paths: list[str]):
+        """Initialize with YAML dataset paths or directories to scan for YOLO datasets."""
+        self.__found_data_file_paths: set[str] = set()
+        self.__data_sets: list[YoloDataFileReader] = []
+
+        self.add_data_sets(data_set_paths)
+        self.SUPPORTED_IMAGE_FORMATS = [
+            ".avif",
+            ".bmp",
+            ".dng",
+            ".heic",
+            ".jp2",
+            ".jpeg",
+            ".jpg",
+            ".mpo",
+            ".png",
+            ".tif",
+            ".tiff",
+            ".webp",
+        ]
+
+    def add_data_sets(self, data_set_paths: list[str]):
+        """Discover YOLO data.yaml files from explicit paths or directory trees."""
+        data_set_paths = [_normalize_path(path) for path in data_set_paths]
+
+        self.__found_data_file_paths.update(
+            [path for path in data_set_paths if os.path.basename(path).lower() == "data.yaml"]
+        )
+
+        paths_without_file_ending = [
+            path for path in data_set_paths if os.path.basename(path).lower() != "data.yaml"
+        ]
+        for path in paths_without_file_ending:
+            self.__found_data_file_paths.update(glob.glob(os.path.join(path, "**", "data.yaml"), recursive=True))
+
+        for path in self.__found_data_file_paths:
+            try:
+                ydfr = YoloDataFileReader(path)
+                self.__data_sets.append(ydfr)
+            except (yaml.YAMLError, TypeError, ValueError) as e:
+                print(f"ERROR: Issue loading YAML file at {path}: {e}")
+            except FileNotFoundError:
+                print(f"ERROR: Yaml file does not exist at {path}")
+
+    def get_found_data_file_paths(self) -> set[str]:
+        """Return all discovered YOLO dataset YAML file paths."""
+        return self.__found_data_file_paths
+
+    def get_found_class_names(self) -> set[str]:
+        """Return the union of all class names from discovered datasets."""
+        found_classes: set[str] = set()
+        for ydf in self.__data_sets:
+            found_classes.update(ydf.get_class_names())
+        return found_classes
+
+    def create_new_dataset_for_class_names(
+        self,
+        class_names: set[str],
+        dst_path: str = ".",
+        data_set_name: str = "new_dataset",
+        ignore_img_formats: bool = False,
+    ) -> bool:
+        """Create a filtered dataset containing only the requested classes."""
+        if len(class_names) == 0:
+            print("ERROR: No classes selected for new dataset.")
+            return False
+
+        new_dataset_path = os.path.join(_normalize_path(dst_path), data_set_name)
+        new_dataset_path = _get_unique_path(new_dataset_path)
+
+        new_class_indices = self.__generate_new_class_indices(class_names)
+        print(f"INFO: New class indices: {new_class_indices}")
+
+        for data_set in self.__data_sets:
+            common_class_names = set(data_set.get_class_names()) & class_names
+            if not common_class_names:
+                continue
+
+            old_to_new_index_match: dict[str, str] = {}
+            for class_name in common_class_names:
+                old_indices = data_set.get_indices_for_names({class_name})
+                if not old_indices:
+                    continue
+                old_index = next(iter(old_indices))
+                new_index = new_class_indices[class_name]
+                old_to_new_index_match[old_index] = new_index
+
+            for split_name in _SPLIT_NAMES:
+                for split_path in data_set.get_split_paths_for_split_name(split_name):
+                    old_images_path = os.path.join(_normalize_path(data_set.get_file_dir()), split_path)
+                    old_labels_path = self.__generate_label_path_from_img_path(old_images_path)
+
+                    missing_path = next(
+                        (path for path in (old_labels_path, old_images_path) if not os.path.exists(path)), None
+                    )
+                    if missing_path is not None:
+                        print(f"WARNING: Path '{missing_path}' does not exist! Skipping it.")
+                        continue
+
+                    new_labels_path = os.path.join(new_dataset_path, split_name, "labels")
+                    new_images_path = os.path.join(new_dataset_path, split_name, "images")
+                    os.makedirs(new_labels_path, exist_ok=True)
+                    os.makedirs(new_images_path, exist_ok=True)
+
+                    for original_label_file_path in glob.glob(os.path.join(old_labels_path, "*.txt")):
+                        label_file = LabelFile(original_label_file_path)
+                        new_label_file_path = label_file.copy_by_class_indices(old_to_new_index_match, new_labels_path)
+                        if new_label_file_path:
+                            file_name = os.path.basename(original_label_file_path)
+                            file_name_no_ext = file_name[: file_name.rfind(".")]
+                            corresponding_image_files = glob.glob(os.path.join(old_images_path, f"{file_name_no_ext}.*"))
+                            for img in corresponding_image_files:
+                                try:
+                                    _, ext = os.path.splitext(img)
+                                except Exception:
+                                    print(f"ERROR: Not a valid image, extension missing for: {img}")
+                                    return False
+
+                                if not ignore_img_formats and ext.lower() not in self.SUPPORTED_IMAGE_FORMATS:
+                                    print(
+                                        f"ERROR: Image format not supported by yolo: {img}\n"
+                                        f"Supported formats are: {self.SUPPORTED_IMAGE_FORMATS}"
+                                    )
+                                    return False
+
+                                new_image_name, _ = os.path.splitext(os.path.basename(new_label_file_path))
+                                new_image_name = new_image_name + ext
+                                destination = os.path.join(new_images_path, new_image_name)
+                                try:
+                                    shutil.copyfile(img, destination)
+                                except OSError as e:
+                                    print(
+                                        f"ERROR: Image could not be copied:\n"
+                                        f"Source: {img}\nDestination: {destination}\n"
+                                        f"Exception was: {e}"
+                                    )
+                                    return False
+
+        data_file_writer = YoloDataFileWriter(new_dataset_path, list(new_class_indices.keys()))
+        try:
+            data_file_writer.write()
+        except Exception as e:
+            print(f"ERROR: Could not write new data.yaml to {new_dataset_path}. Generation of {data_set_name} is incomplete!")
+            print(f"Exception was: {e}")
+            return False
+
+        return True
+
+    def __generate_new_class_indices(self, class_names: set[str]) -> dict[str, str]:
+        """Generate a fresh label index mapping for the selected class names."""
+        ret: dict[str, str] = {}
+        for i, name in enumerate(class_names):
+            ret[name] = str(i)
+        return ret
+
+    def __generate_label_path_from_img_path(self, img_path: str) -> str:
+        """Convert an images directory path into its corresponding labels directory path."""
+        img_path = _normalize_path(img_path)
+        return os.path.normpath(os.path.join(os.path.dirname(img_path), "labels"))

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/labels.py

@@ -0,0 +1,71 @@
+import os
+import re
+
+from .paths import _get_unique_path
+
+
+class LabelFile:
+    """Read YOLO label files and copy label lines with remapped class indices."""
+
+    def __init__(self, file_path: str):
+        """
+        Args:
+            file_path: Valid path to the label file
+        Raises:
+            FileNotFoundError: When the label file does not exist
+        """
+        self.__class_indices: set[str] = set()
+        self.file_path = _get_unique_path(file_path) if os.path.isdir(file_path) else os.path.abspath(file_path)
+        if not os.path.exists(self.file_path) or not os.path.isfile(self.file_path):
+            raise FileNotFoundError(f"No label file found at: {self.file_path}")
+
+        self.read_label_file()
+
+    def read_label_file(self):
+        """Parse the label file and collect all class indices referenced in it."""
+        with open(self.file_path, encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                match = re.match(r"^(\d+) ", line)
+                if not match:
+                    print(f"WARNING: No class found on line {i} in label file {self.file_path}")
+                else:
+                    self.__class_indices.add(match.group(1))
+
+    def get_class_indices(self) -> set[str]:
+        """Return the set of class indices in the loaded label file."""
+        return self.__class_indices
+
+    def copy_by_class_indices(self, old_to_new_index_match: dict[str, str], dst_path: str, file_name: str = "") -> str | None:
+        """Copy the label file to the destination with indices remapped for the new dataset."""
+        common_class_indices = self.__class_indices & set(old_to_new_index_match.keys())
+        if not common_class_indices:
+            return None
+
+        dst_path = os.path.abspath(dst_path)
+        if os.path.basename(dst_path) != "labels":
+            dst_path = os.path.join(dst_path, "labels")
+
+        if os.path.exists(dst_path) and os.path.isdir(dst_path):
+            new_content: list[str] = []
+            with open(self.file_path, encoding="utf-8") as f:
+                for line in f:
+                    match = re.match(r"^(\d+) ", line)
+                    if not match:
+                        continue
+                    elif match.group(1) in common_class_indices:
+                        old_index = match.group(1)
+                        new_index = old_to_new_index_match[old_index]
+                        if new_index != old_index:
+                            line = line.replace(old_index, new_index, 1)
+                        new_content.append(line)
+
+            if file_name == "":
+                file_name = os.path.basename(self.file_path)
+            dst_file_path = _get_unique_path(os.path.join(dst_path, file_name))
+
+            with open(dst_file_path, "w", encoding="utf-8") as f:
+                f.writelines(new_content)
+        else:
+            raise FileNotFoundError(f"Directory not found: {dst_path}")
+
+        return dst_file_path

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/paths.py

@@ -0,0 +1,20 @@
+import os
+
+_SPLIT_NAMES = ["train", "val", "test"]
+
+
+def _normalize_path(path: str) -> str:
+    """Normalize path separators and return an absolute path for cross-platform use."""
+    return os.path.abspath(str(path).replace("\\", os.sep))
+
+
+def _get_unique_path(path: str) -> str:
+    """Return an unused file path by appending a numeric suffix when needed."""
+    path = _normalize_path(path)
+    base, ext = os.path.splitext(path)
+    unique_path = path
+    index = 1
+    while os.path.exists(unique_path):
+        unique_path = f"{base}_{index}{ext}"
+        index += 1
+    return unique_path

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/reader.py

@@ -0,0 +1,128 @@
+import os
+from typing import Any
+
+import yaml
+
+from .paths import _SPLIT_NAMES, _normalize_path
+
+
+class _UniqueKeyLoader(yaml.SafeLoader):
+    """A YAML SafeLoader that rejects mappings with duplicate keys.
+
+    PyYAML's default loaders silently keep only the last value for a duplicate
+    mapping key (e.g. ``0: car`` followed by ``0: person`` collapses to
+    ``{0: 'person'}``)"""
+
+    def construct_mapping(self, node: yaml.MappingNode, deep: bool = False) -> dict[Any, Any]:
+        seen_keys: set[Any] = set()
+        for key_node, _ in node.value:
+            key = self.construct_object(key_node, deep=deep)
+            if key in seen_keys:
+                raise yaml.constructor.ConstructorError(
+                    "while constructing a mapping", node.start_mark, f"found duplicate key: {key!r}", key_node.start_mark
+                )
+            seen_keys.add(key)
+        return super().construct_mapping(node, deep)
+
+
+class YoloDataFileReader:
+    """Read a YOLO dataset YAML file and expose its split paths and class names."""
+
+    def __init__(self, file_path: str):
+        """
+        Args:
+            file_path: Valid path to the data set YAML file
+        Raises:
+            yaml.YAMLError: When there is an issue with reading the data-set YAML file
+            FileNotFoundError: When the data-set YAML file does not exist
+            ValueError: When a value in the YAML file is missing or is faulty
+            TypeError: When a value comes in an unexpected type
+        """
+        self.__file_path = _normalize_path(file_path)
+        self.__file_dir = os.path.dirname(self.__file_path)
+        self.__splits_paths: dict[str, list[str]] = {}
+        self.__name_to_index: dict[str, int] = {}
+
+        self.sync_content()
+
+    def get_split_paths_for_split_name(self, split_name: str) -> list[str]:
+        """Return a list of dataset split paths for the requested split name."""
+        return self.__splits_paths.get(split_name, [])
+
+    def get_class_names(self) -> list[str]:
+        """Return the class names defined in the dataset YAML, ordered by their class index."""
+        return list(self.__name_to_index.keys())
+
+    def get_file_path(self) -> str:
+        """Return the absolute path to the loaded YAML file."""
+        return self.__file_path
+
+    def get_file_dir(self) -> str:
+        """Return the directory containing the loaded dataset YAML file."""
+        return self.__file_dir
+
+    def get_indices_for_names(self, class_names: set[str]) -> set[str]:
+        """Map a set of class names to their corresponding numeric label indices."""
+        return {str(self.__name_to_index[class_name]) for class_name in class_names if class_name in self.__name_to_index}
+
+    def sync_content(self):
+        """Reload YAML content and refresh split and name mappings."""
+        with open(self.__file_path, encoding="utf-8") as f:
+            content = yaml.load(f, Loader=_UniqueKeyLoader)
+        for split_name in _SPLIT_NAMES:
+            self.__splits_paths[split_name] = self.__read_yaml_entry(content, split_name)
+        self.__name_to_index = self.__read_class_names_entry(content)
+
+    def __read_class_names_entry(self, yaml_content: Any) -> dict[str, int]:
+        """Parse the 'names' entry into a name-to-index map, ordered by class index."""
+        entry = yaml_content.get("names")
+        if not entry:
+            return {}
+
+        if isinstance(entry, str):
+            return {entry: 0}
+        elif isinstance(entry, dict):
+            try:
+                index_to_name = {int(index): name for index, name in entry.items()}
+            except (TypeError, ValueError) as e:
+                raise TypeError(f"Class indices in 'names' must be integers in: {self.__file_path}") from e
+            if len(index_to_name) != len(entry):
+                raise ValueError(f"Duplicate class indices found in 'names' in: {self.__file_path}")
+            if len(set(index_to_name.values())) != len(index_to_name):
+                raise ValueError(f"Duplicate class names found in 'names' in: {self.__file_path}")
+            ordered_indices = sorted(index_to_name)
+            return {index_to_name[index]: index for index in ordered_indices}
+        elif isinstance(entry, list):
+            if len(set(entry)) != len(entry):
+                raise ValueError(f"Duplicate class names found in 'names' in: {self.__file_path}")
+            return {name: index for index, name in enumerate(entry)}
+        else:
+            raise TypeError(f"Unexpected type for names in: {self.__file_path}")
+
+    def __read_yaml_entry(self, yaml_content: Any, entry_name: str, missing_ok: bool = True) -> list[str]:
+        """Normalize YAML entry values into a list of paths or names."""
+        entry = yaml_content.get(entry_name)
+        ret: list[str] = []
+        if entry:
+            if isinstance(entry, str):
+                ret.append(entry)
+            elif isinstance(entry, dict):
+                ret.extend(entry.values())
+            elif isinstance(entry, list):
+                ret.extend(entry)
+            else:
+                raise TypeError(f"Unexpected type for {entry_name} in: {self.__file_path}")
+        else:
+            if not missing_ok:
+                raise ValueError(f"{entry_name} not defined in: {self.__file_path}")
+
+        return ret
+
+    def __repr__(self) -> str:
+        return (
+            f"\nYoloDatasetFile(\n"
+            f"    path: {self.__file_path}\n"
+            f"    split_paths: {self.__splits_paths}\n"
+            f"    class_names: {self.get_class_names()}\n"
+            f")\n"
+        )

yolodatasetcustomizer-0.3.0/src/YoloDatasetCustomizer/writer.py

@@ -0,0 +1,41 @@
+import os
+
+import yaml
+
+from .paths import _normalize_path
+
+
+class YoloDataFileWriter:
+    """Create and write a YOLO dataset YAML file for a generated dataset."""
+
+    def __init__(
+        self,
+        data_set_dir: str = "./",
+        class_names: list[str] | None = None,
+        train_split: str = "./train",
+        val_split: str = "./valid",
+        test_split: str = "./test",
+    ) -> None:
+        """Initialize writer with dataset directory, class names, and split locations."""
+        self.data_set_dir = _normalize_path(data_set_dir)
+        self.class_names = list(class_names) if class_names is not None else []
+        self.train_split = train_split
+        self.val_split = val_split
+        self.test_split = test_split
+        self.FILE_NAME = "data.yaml"
+
+    def write(self):
+        """Write the YOLO data.yaml file with train/val/test paths and class names."""
+        classes = {}
+        for i, class_name in enumerate(self.class_names):
+            classes[i] = class_name
+        data = {
+            "train": f"{self.train_split}/images",
+            "val": f"{self.val_split}/images",
+            "test": f"{self.test_split}/images",
+            "names": classes,
+            "nc": len(self.class_names),
+        }
+
+        with open(os.path.join(self.data_set_dir, self.FILE_NAME), "w", encoding="utf-8") as f:
+            yaml.safe_dump(data, f, default_flow_style=False)