YoloDatasetCustomizer 0.2.0__tar.gz → 0.4.0__tar.gz

yolodatasetcustomizer-0.4.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.2.0 → yolodatasetcustomizer-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: YoloDatasetCustomizer
-Version: 0.2.0
+Version: 0.4.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
@@ -23,6 +23,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Image Recognition
 Requires-Python: >=3.10
 Requires-Dist: pyyaml>=6.0.3
+Requires-Dist: rich>=13.0.0
 Description-Content-Type: text/markdown
 
 # Intro
@@ -42,15 +43,21 @@ If a path ends on _data.yaml_ it will take that file into account.
 If a path doesn't end on _/data.yaml_ it will recursevly look for all _data.yaml_ files from that directory downards.
 e.g.
 
-    from YoloDatasetCustomizer import YoloDatasetCustomizer
+```python
+from YoloDatasetCustomizer import YoloDatasetCustomizer
 
-    ydc = YoloDatasetCustomizer(list_to_existing_datasets)
-    ydc.add([one_more_existing_dataset])
+ydc = YoloDatasetCustomizer(list_to_existing_datasets)
+ydc.add_data_sets([one_more_existing_dataset])
 
-    if ydc.create_new_dataset_for_class_names(set_of_wanted_class_names, optional_custom_path, optional_custom_name):
-        print("All good")
-    else:
-        print("Oh no.")
+if ydc.create_new_dataset_for_class_names(
+    class_names=<set_of_wanted_class_names>,
+    dst_path=<optional_custom_path>,
+    data_set_name=<optional_custom_name>,
+):
+    print("All good")
+else:
+    print("Oh no.")
+```
 
 # Run tests
 
@@ -58,13 +65,18 @@ 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
+Build the package:
 
     uv build
 
-Publish to PyPI (still need twine)
+Publish to PyPI (still need twine):
 
     uv pip install twine
     uv run twine upload dist/\*

{yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.0}/README.md

@@ -15,15 +15,21 @@ If a path ends on _data.yaml_ it will take that file into account.
 If a path doesn't end on _/data.yaml_ it will recursevly look for all _data.yaml_ files from that directory downards.
 e.g.
 
-    from YoloDatasetCustomizer import YoloDatasetCustomizer
-
-    ydc = YoloDatasetCustomizer(list_to_existing_datasets)
-    ydc.add([one_more_existing_dataset])
-
-    if ydc.create_new_dataset_for_class_names(set_of_wanted_class_names, optional_custom_path, optional_custom_name):
-        print("All good")
-    else:
-        print("Oh no.")
+```python
+from YoloDatasetCustomizer import YoloDatasetCustomizer
+
+ydc = YoloDatasetCustomizer(list_to_existing_datasets)
+ydc.add_data_sets([one_more_existing_dataset])
+
+if ydc.create_new_dataset_for_class_names(
+    class_names=<set_of_wanted_class_names>,
+    dst_path=<optional_custom_path>,
+    data_set_name=<optional_custom_name>,
+):
+    print("All good")
+else:
+    print("Oh no.")
+```
 
 # Run tests
 
@@ -31,13 +37,18 @@ 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
+Build the package:
 
     uv build
 
-Publish to PyPI (still need twine)
+Publish to PyPI (still need twine):
 
     uv pip install twine
     uv run twine upload dist/\*

{yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "YoloDatasetCustomizer"
-version = "0.2.0"
+version = "0.4.0"
 description = "Load one or multiple datasets. Reduce, mix or merge into a custom dataset."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -30,6 +30,7 @@ classifiers = [
 
 dependencies = [
     "pyyaml>=6.0.3",
+    "rich>=13.0.0",
 ]
 
 [project.urls]
@@ -41,4 +42,18 @@ Repository = "https://github.com/mLuca/YoloDatasetCustomizer"
 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.4.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.2.0 → yolodatasetcustomizer-0.4.0}/src/YoloDatasetCustomizer/__init__.py

@@ -5,7 +5,7 @@ from .labels import LabelFile
 from .reader import YoloDataFileReader
 from .writer import YoloDataFileWriter
 
-__version__ = "0.2.0"
+__version__ = "0.4.0"
 __author__ = "Luca Mazzon"
 __email__ = "mazzon.luca@gmail.com"
 

yolodatasetcustomizer-0.4.0/src/YoloDatasetCustomizer/customizer.py

@@ -0,0 +1,223 @@
+import glob
+import os
+import shutil
+
+import yaml
+from rich.progress import BarColumn, MofNCompleteColumn, Progress, TaskProgressColumn, TextColumn
+
+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}")
+
+        # Always create every split folder so the folder structure matches what data.yaml references,
+        # even when no source dataset contributed images to a split.
+        for split_name in _SPLIT_NAMES:
+            os.makedirs(os.path.join(new_dataset_path, split_name, "labels"), exist_ok=True)
+            os.makedirs(os.path.join(new_dataset_path, split_name, "images"), exist_ok=True)
+
+        # Plan first, so every stage knows how many files it is going to copy
+        planned_files: dict[str, list[tuple[LabelFile, dict[str, list[str]], dict[str, str]]]] = {
+            split_name: [] for split_name in _SPLIT_NAMES
+        }
+        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
+
+                    # Index images by exact stem so 'car.*' can't pick up e.g. 'car.v2.jpg'
+                    images_by_stem: dict[str, list[str]] = {}
+                    for image_name in os.listdir(old_images_path):
+                        stem = os.path.splitext(image_name)[0]
+                        images_by_stem.setdefault(stem, []).append(os.path.join(old_images_path, image_name))
+
+                    for original_label_file_path in glob.glob(os.path.join(old_labels_path, "*.txt")):
+                        label_file = LabelFile(original_label_file_path)
+                        if label_file.get_class_indices() & set(old_to_new_index_match):
+                            planned_files[split_name].append((label_file, images_by_stem, old_to_new_index_match))
+
+        # One bar per stage (white, filled green), all visible from the start
+        with Progress(
+            TextColumn("{task.description}"),
+            BarColumn(style="white", complete_style="green", finished_style="green"),
+            TaskProgressColumn(),
+            MofNCompleteColumn(),
+        ) as progress:
+            stage_tasks = {
+                split_name: progress.add_task(f"[{stage}/{len(_SPLIT_NAMES)}] {split_name:<5}", total=len(planned_files[split_name]))
+                for stage, split_name in enumerate(_SPLIT_NAMES, start=1)
+            }
+            for split_name in _SPLIT_NAMES:
+                new_labels_path = os.path.join(new_dataset_path, split_name, "labels")
+                new_images_path = os.path.join(new_dataset_path, split_name, "images")
+                for label_file, images_by_stem, old_to_new_index_match in planned_files[split_name]:
+                    if not self.__copy_label_and_images(
+                        label_file, images_by_stem, old_to_new_index_match, new_labels_path, new_images_path, ignore_img_formats
+                    ):
+                        return False
+                    progress.advance(stage_tasks[split_name])
+
+        # Pass the split folder names explicitly so data.yaml points at the folders created above
+        data_file_writer = YoloDataFileWriter(
+            new_dataset_path,
+            list(new_class_indices.keys()),
+            train_split="train",
+            val_split="val",
+            test_split="test",
+        )
+        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 __copy_label_and_images(
+        self,
+        label_file: LabelFile,
+        images_by_stem: dict[str, list[str]],
+        old_to_new_index_match: dict[str, str],
+        new_labels_path: str,
+        new_images_path: str,
+        ignore_img_formats: bool,
+    ) -> bool:
+        """Copy one label file with remapped indices plus its images. Returns False on error."""
+        new_label_file_path = label_file.copy_by_class_indices(old_to_new_index_match, new_labels_path)
+        if not new_label_file_path:
+            return True
+
+        file_name_no_ext = os.path.splitext(os.path.basename(label_file.file_path))[0]
+        for img in images_by_stem.get(file_name_no_ext, []):
+            _, ext = os.path.splitext(img)
+
+            if not ignore_img_formats and ext.lower() not in self.SUPPORTED_IMAGE_FORMATS:
+                print(
+                    f"\nERROR: 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))
+            destination = os.path.join(new_images_path, new_image_name + ext)
+            try:
+                shutil.copyfile(img, destination)
+            except OSError as e:
+                print(
+                    f"\nERROR: Image could not be copied:\n"
+                    f"Source: {img}\nDestination: {destination}\n"
+                    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.2.0 → yolodatasetcustomizer-0.4.0}/src/YoloDatasetCustomizer/labels.py

@@ -1,71 +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_indeces: 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_indeces.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_indeces
-
-    def copy_by_class_indeces(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_indeces = self.__class_indeces & set(old_to_new_index_match.keys())
-        if not common_class_indeces:
-            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_indeces:
-                        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
+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.2.0 → yolodatasetcustomizer-0.4.0}/src/YoloDatasetCustomizer/paths.py

@@ -1,20 +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
+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