YoloDatasetCustomizer 0.2.0__tar.gz → 0.4.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (19) hide show
  1. yolodatasetcustomizer-0.4.1/.github/workflows/ci.yml +35 -0
  2. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/PKG-INFO +22 -10
  3. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/README.md +22 -11
  4. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/pyproject.toml +59 -44
  5. yolodatasetcustomizer-0.4.1/resources/Improvements.txt +57 -0
  6. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/src/YoloDatasetCustomizer/__init__.py +1 -1
  7. yolodatasetcustomizer-0.4.1/src/YoloDatasetCustomizer/customizer.py +235 -0
  8. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/src/YoloDatasetCustomizer/labels.py +71 -71
  9. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/src/YoloDatasetCustomizer/paths.py +20 -20
  10. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/src/YoloDatasetCustomizer/reader.py +128 -90
  11. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/src/YoloDatasetCustomizer/writer.py +41 -41
  12. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/test/YoloDatasetCustomizer_test.py +167 -7
  13. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/uv.lock +1117 -808
  14. yolodatasetcustomizer-0.2.0/resources/Improvements.txt +0 -33
  15. yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/customizer.py +0 -166
  16. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/.gitattributes +0 -0
  17. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/.gitignore +0 -0
  18. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/.python-version +0 -0
  19. {yolodatasetcustomizer-0.2.0 → yolodatasetcustomizer-0.4.1}/LICENSE +0 -0
@@ -0,0 +1,35 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ubuntu-latest
11
+ strategy:
12
+ fail-fast: false
13
+ matrix:
14
+ # "3" resolves to the latest released stable CPython 3.x, so this
15
+ # tracks new releases automatically without needing manual bumps
16
+ python-version: ["3.10", "3"]
17
+ steps:
18
+ - uses: actions/checkout@v4
19
+
20
+ - name: Install uv
21
+ uses: astral-sh/setup-uv@v6
22
+ with:
23
+ python-version: ${{ matrix.python-version }}
24
+
25
+ - name: Install dependencies
26
+ run: uv sync --group dev
27
+
28
+ - name: Lint (ruff)
29
+ run: uv run ruff check src test
30
+
31
+ - name: Type-check (mypy)
32
+ run: uv run mypy
33
+
34
+ - name: Test (pytest)
35
+ run: uv run pytest -q
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: YoloDatasetCustomizer
3
- Version: 0.2.0
3
+ Version: 0.4.1
4
4
  Summary: Load one or multiple datasets. Reduce, mix or merge into a custom dataset.
5
5
  Project-URL: Homepage, https://github.com/mLuca/YoloDatasetCustomizer
6
6
  Project-URL: Repository, https://github.com/mLuca/YoloDatasetCustomizer
@@ -23,6 +23,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
23
23
  Classifier: Topic :: Scientific/Engineering :: Image Recognition
24
24
  Requires-Python: >=3.10
25
25
  Requires-Dist: pyyaml>=6.0.3
26
+ Requires-Dist: rich>=13.0.0
26
27
  Description-Content-Type: text/markdown
27
28
 
28
29
  # Intro
@@ -42,15 +43,21 @@ If a path ends on _data.yaml_ it will take that file into account.
42
43
  If a path doesn't end on _/data.yaml_ it will recursevly look for all _data.yaml_ files from that directory downards.
43
44
  e.g.
44
45
 
45
- from YoloDatasetCustomizer import YoloDatasetCustomizer
46
+ ```python
47
+ from YoloDatasetCustomizer import YoloDatasetCustomizer
46
48
 
47
- ydc = YoloDatasetCustomizer(list_to_existing_datasets)
48
- ydc.add([one_more_existing_dataset])
49
+ ydc = YoloDatasetCustomizer(list_to_existing_datasets)
50
+ ydc.add_data_sets([one_more_existing_dataset])
49
51
 
50
- if ydc.create_new_dataset_for_class_names(set_of_wanted_class_names, optional_custom_path, optional_custom_name):
51
- print("All good")
52
- else:
53
- print("Oh no.")
52
+ if ydc.create_new_dataset_for_class_names(
53
+ class_names=<set_of_wanted_class_names>,
54
+ dst_path=<optional_custom_path>,
55
+ data_set_name=<optional_custom_name>,
56
+ ):
57
+ print("All good")
58
+ else:
59
+ print("Oh no.")
60
+ ```
54
61
 
55
62
  # Run tests
56
63
 
@@ -58,13 +65,18 @@ The project uses uv. Install dependencies and run tests with:
58
65
 
59
66
  uv run pytest
60
67
 
68
+ # Lint & type-check
69
+
70
+ uv run ruff check src test
71
+ uv run mypy
72
+
61
73
  # Build project
62
74
 
63
- Build the package
75
+ Build the package:
64
76
 
65
77
  uv build
66
78
 
67
- Publish to PyPI (still need twine)
79
+ Publish to PyPI (still need twine):
68
80
 
69
81
  uv pip install twine
70
82
  uv run twine upload dist/\*
@@ -15,15 +15,21 @@ If a path ends on _data.yaml_ it will take that file into account.
15
15
  If a path doesn't end on _/data.yaml_ it will recursevly look for all _data.yaml_ files from that directory downards.
16
16
  e.g.
17
17
 
18
- from YoloDatasetCustomizer import YoloDatasetCustomizer
19
-
20
- ydc = YoloDatasetCustomizer(list_to_existing_datasets)
21
- ydc.add([one_more_existing_dataset])
22
-
23
- if ydc.create_new_dataset_for_class_names(set_of_wanted_class_names, optional_custom_path, optional_custom_name):
24
- print("All good")
25
- else:
26
- print("Oh no.")
18
+ ```python
19
+ from YoloDatasetCustomizer import YoloDatasetCustomizer
20
+
21
+ ydc = YoloDatasetCustomizer(list_to_existing_datasets)
22
+ ydc.add_data_sets([one_more_existing_dataset])
23
+
24
+ if ydc.create_new_dataset_for_class_names(
25
+ class_names=<set_of_wanted_class_names>,
26
+ dst_path=<optional_custom_path>,
27
+ data_set_name=<optional_custom_name>,
28
+ ):
29
+ print("All good")
30
+ else:
31
+ print("Oh no.")
32
+ ```
27
33
 
28
34
  # Run tests
29
35
 
@@ -31,13 +37,18 @@ The project uses uv. Install dependencies and run tests with:
31
37
 
32
38
  uv run pytest
33
39
 
40
+ # Lint & type-check
41
+
42
+ uv run ruff check src test
43
+ uv run mypy
44
+
34
45
  # Build project
35
46
 
36
- Build the package
47
+ Build the package:
37
48
 
38
49
  uv build
39
50
 
40
- Publish to PyPI (still need twine)
51
+ Publish to PyPI (still need twine):
41
52
 
42
53
  uv pip install twine
43
54
  uv run twine upload dist/\*
@@ -1,44 +1,59 @@
1
- [build-system]
2
- requires = ["hatchling"]
3
- build-backend = "hatchling.build"
4
-
5
- [project]
6
- name = "YoloDatasetCustomizer"
7
- version = "0.2.0"
8
- description = "Load one or multiple datasets. Reduce, mix or merge into a custom dataset."
9
- readme = "README.md"
10
- requires-python = ">=3.10"
11
- license = {text = "GPL-3.0-only"}
12
- authors = [
13
- {name = "Luca Mazzon", email = "mazzon.luca@gmail.com"},
14
- ]
15
- keywords = ["ultralytics", "yolo", "dataset", "machine-learning", "computer-vision", "customization"]
16
- classifiers = [
17
- "Development Status :: 5 - Production/Stable",
18
- "Intended Audience :: Developers",
19
- "Intended Audience :: Science/Research",
20
- "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
21
- "Programming Language :: Python :: 3",
22
- "Programming Language :: Python :: 3.10",
23
- "Programming Language :: Python :: 3.11",
24
- "Programming Language :: Python :: 3.12",
25
- "Programming Language :: Python :: 3.13",
26
- "Programming Language :: Python :: 3.14",
27
- "Topic :: Scientific/Engineering :: Artificial Intelligence",
28
- "Topic :: Scientific/Engineering :: Image Recognition",
29
- ]
30
-
31
- dependencies = [
32
- "pyyaml>=6.0.3",
33
- ]
34
-
35
- [project.urls]
36
- Homepage = "https://github.com/mLuca/YoloDatasetCustomizer"
37
- Repository = "https://github.com/mLuca/YoloDatasetCustomizer"
38
- "Issue Tracker" = "https://github.com/mLuca/YoloDatasetCustomizer/issues"
39
-
40
- [dependency-groups]
41
- dev = [
42
- "pytest>=9.1.1",
43
- "twine>=6.2.0",
44
- ]
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "YoloDatasetCustomizer"
7
+ version = "0.4.1"
8
+ description = "Load one or multiple datasets. Reduce, mix or merge into a custom dataset."
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = {text = "GPL-3.0-only"}
12
+ authors = [
13
+ {name = "Luca Mazzon", email = "mazzon.luca@gmail.com"},
14
+ ]
15
+ keywords = ["ultralytics", "yolo", "dataset", "machine-learning", "computer-vision", "customization"]
16
+ classifiers = [
17
+ "Development Status :: 5 - Production/Stable",
18
+ "Intended Audience :: Developers",
19
+ "Intended Audience :: Science/Research",
20
+ "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
21
+ "Programming Language :: Python :: 3",
22
+ "Programming Language :: Python :: 3.10",
23
+ "Programming Language :: Python :: 3.11",
24
+ "Programming Language :: Python :: 3.12",
25
+ "Programming Language :: Python :: 3.13",
26
+ "Programming Language :: Python :: 3.14",
27
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
28
+ "Topic :: Scientific/Engineering :: Image Recognition",
29
+ ]
30
+
31
+ dependencies = [
32
+ "pyyaml>=6.0.3",
33
+ "rich>=13.0.0",
34
+ ]
35
+
36
+ [project.urls]
37
+ Homepage = "https://github.com/mLuca/YoloDatasetCustomizer"
38
+ Repository = "https://github.com/mLuca/YoloDatasetCustomizer"
39
+ "Issue Tracker" = "https://github.com/mLuca/YoloDatasetCustomizer/issues"
40
+
41
+ [dependency-groups]
42
+ dev = [
43
+ "pytest>=9.1.1",
44
+ "twine>=6.2.0",
45
+ "ruff>=0.14.0",
46
+ "mypy>=1.18.0",
47
+ ]
48
+
49
+ [tool.ruff]
50
+ line-length = 120
51
+
52
+ [tool.ruff.lint]
53
+ select = ["E", "F", "W", "I"]
54
+ ignore = ["E501"]
55
+
56
+ [tool.mypy]
57
+ python_version = "3.10"
58
+ files = ["src"]
59
+ ignore_missing_imports = true
@@ -0,0 +1,57 @@
1
+ Priority-ordered improvement list (last reviewed 2026-08-14)
2
+
3
+ FIXED 1. Fix brittle path and OS handling
4
+ 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.
5
+
6
+ FIXED 2. Remove mutable default arguments
7
+ 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.
8
+
9
+ FIXED 3. Fix broken internal state and debug output
10
+ 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.
11
+
12
+ FIXED 4. Refactor the monolithic implementation into smaller modules
13
+ 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.
14
+
15
+ FIXED 1. Fix broken copy-failure error handling in customizer.py
16
+ 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.
17
+
18
+ FIXED 2. Fix the path-existence skip logic in create_new_dataset_for_class_names
19
+ 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).
20
+
21
+ FIXED 3. Standardize class-index typing and conversion
22
+ 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.
23
+
24
+ 4. Tighten validation and error handling generally
25
+ 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.
26
+
27
+ 5. Stop leaking mutable internal state from getters
28
+ 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.
29
+
30
+ FIXED 6. Add CI (GitHub Actions) for tests, lint, and type-checking
31
+ 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.
32
+
33
+ FIXED 7. Add lint/type-check tooling (ruff + mypy) to the dev dependency group
34
+ 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.
35
+
36
+ 8. Replace print()-based warnings/errors with the `logging` module
37
+ Why: A library shouldn't write directly to stdout. Using `logging` lets consumers control verbosity, redirect output, and integrate the tool into larger pipelines.
38
+
39
+ FIXED 9. Clean up naming and spelling inconsistencies
40
+ 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.
41
+
42
+ 10. Remove or clarify the directory-handling branch in LabelFile.__init__
43
+ 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.
44
+
45
+ 11. Make SUPPORTED_IMAGE_FORMATS a class-level constant
46
+ 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.
47
+
48
+ 12. Add a py.typed marker file
49
+ 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.
50
+
51
+ 13. Improve test coverage around edge cases
52
+ 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.
53
+
54
+ 14. Clarify and formally validate documented limitations
55
+ 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.
56
+
57
+ 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.
@@ -5,7 +5,7 @@ from .labels import LabelFile
5
5
  from .reader import YoloDataFileReader
6
6
  from .writer import YoloDataFileWriter
7
7
 
8
- __version__ = "0.2.0"
8
+ __version__ = "0.4.1"
9
9
  __author__ = "Luca Mazzon"
10
10
  __email__ = "mazzon.luca@gmail.com"
11
11
 
@@ -0,0 +1,235 @@
1
+ import glob
2
+ import os
3
+ import shutil
4
+
5
+ import yaml
6
+ from rich.progress import BarColumn, MofNCompleteColumn, Progress, TaskProgressColumn, TextColumn
7
+
8
+ from .labels import LabelFile
9
+ from .paths import _SPLIT_NAMES, _get_unique_path, _normalize_path
10
+ from .reader import YoloDataFileReader
11
+ from .writer import YoloDataFileWriter
12
+
13
+
14
+ class YoloDatasetCustomizer:
15
+ """Find and customize YOLO datasets across one or more directory paths."""
16
+
17
+ def __init__(self, data_set_paths: list[str]):
18
+ """Initialize with YAML dataset paths or directories to scan for YOLO datasets."""
19
+ self.__found_data_file_paths: set[str] = set()
20
+ self.__data_sets: list[YoloDataFileReader] = []
21
+
22
+ self.add_data_sets(data_set_paths)
23
+ self.SUPPORTED_IMAGE_FORMATS = [
24
+ ".avif",
25
+ ".bmp",
26
+ ".dng",
27
+ ".heic",
28
+ ".jp2",
29
+ ".jpeg",
30
+ ".jpg",
31
+ ".mpo",
32
+ ".png",
33
+ ".tif",
34
+ ".tiff",
35
+ ".webp",
36
+ ]
37
+
38
+ def add_data_sets(self, data_set_paths: list[str]):
39
+ """Discover YOLO data.yaml files from explicit paths or directory trees."""
40
+ data_set_paths = [_normalize_path(path) for path in data_set_paths]
41
+
42
+ self.__found_data_file_paths.update(
43
+ [path for path in data_set_paths if os.path.basename(path).lower() == "data.yaml"]
44
+ )
45
+
46
+ paths_without_file_ending = [
47
+ path for path in data_set_paths if os.path.basename(path).lower() != "data.yaml"
48
+ ]
49
+ for path in paths_without_file_ending:
50
+ self.__found_data_file_paths.update(glob.glob(os.path.join(path, "**", "data.yaml"), recursive=True))
51
+
52
+ for path in self.__found_data_file_paths:
53
+ try:
54
+ ydfr = YoloDataFileReader(path)
55
+ self.__data_sets.append(ydfr)
56
+ except (yaml.YAMLError, TypeError, ValueError) as e:
57
+ print(f"ERROR: Issue loading YAML file at {path}: {e}")
58
+ except FileNotFoundError:
59
+ print(f"ERROR: Yaml file does not exist at {path}")
60
+
61
+ def get_found_data_file_paths(self) -> set[str]:
62
+ """Return all discovered YOLO dataset YAML file paths."""
63
+ return self.__found_data_file_paths
64
+
65
+ def get_found_class_names(self) -> set[str]:
66
+ """Return the union of all class names from discovered datasets."""
67
+ found_classes: set[str] = set()
68
+ for ydf in self.__data_sets:
69
+ found_classes.update(ydf.get_class_names())
70
+ return found_classes
71
+
72
+ def create_new_dataset_for_class_names(
73
+ self,
74
+ class_names: set[str],
75
+ dst_path: str = ".",
76
+ data_set_name: str = "new_dataset",
77
+ ignore_img_formats: bool = False,
78
+ ) -> bool:
79
+ """Create a filtered dataset containing only the requested classes."""
80
+ if len(class_names) == 0:
81
+ print("ERROR: No classes selected for new dataset.")
82
+ return False
83
+
84
+ new_dataset_path = os.path.join(_normalize_path(dst_path), data_set_name)
85
+ new_dataset_path = _get_unique_path(new_dataset_path)
86
+
87
+ new_class_indices = self.__generate_new_class_indices(class_names)
88
+ print(f"INFO: New class indices: {new_class_indices}")
89
+
90
+ # Always create every split folder so the folder structure matches what data.yaml references,
91
+ # even when no source dataset contributed images to a split.
92
+ for split_name in _SPLIT_NAMES:
93
+ os.makedirs(os.path.join(new_dataset_path, split_name, "labels"), exist_ok=True)
94
+ os.makedirs(os.path.join(new_dataset_path, split_name, "images"), exist_ok=True)
95
+
96
+ # Plan first, so every stage knows how many files it is going to copy
97
+ planned_files: dict[str, list[tuple[LabelFile, dict[str, list[str]], dict[str, str]]]] = {
98
+ split_name: [] for split_name in _SPLIT_NAMES
99
+ }
100
+ for data_set in self.__data_sets:
101
+ common_class_names = set(data_set.get_class_names()) & class_names
102
+ if not common_class_names:
103
+ continue
104
+
105
+ old_to_new_index_match: dict[str, str] = {}
106
+ for class_name in common_class_names:
107
+ old_indices = data_set.get_indices_for_names({class_name})
108
+ if not old_indices:
109
+ continue
110
+ old_index = next(iter(old_indices))
111
+ new_index = new_class_indices[class_name]
112
+ old_to_new_index_match[old_index] = new_index
113
+
114
+ for split_name in _SPLIT_NAMES:
115
+ for split_path in data_set.get_split_paths_for_split_name(split_name):
116
+ old_images_path = os.path.join(_normalize_path(data_set.get_file_dir()), split_path)
117
+ old_labels_path = self.__generate_label_path_from_img_path(old_images_path)
118
+
119
+ missing_path = next(
120
+ (path for path in (old_labels_path, old_images_path) if not os.path.exists(path)), None
121
+ )
122
+ if missing_path is not None:
123
+ print(f"WARNING: Path '{missing_path}' does not exist! Skipping it.")
124
+ continue
125
+
126
+ # Index images by exact stem so 'car.*' can't pick up e.g. 'car.v2.jpg'.
127
+ # Only regular files with an extension count, like the old 'car.*' lookup.
128
+ images_by_stem: dict[str, list[str]] = {}
129
+ for image_name in os.listdir(old_images_path):
130
+ stem, ext = os.path.splitext(image_name)
131
+ image_path = os.path.join(old_images_path, image_name)
132
+ if not ext or not os.path.isfile(image_path):
133
+ continue
134
+ images_by_stem.setdefault(stem, []).append(image_path)
135
+
136
+ for original_label_file_path in glob.glob(os.path.join(old_labels_path, "*.txt")):
137
+ label_file = LabelFile(original_label_file_path)
138
+ if label_file.get_class_indices() & set(old_to_new_index_match):
139
+ planned_files[split_name].append((label_file, images_by_stem, old_to_new_index_match))
140
+
141
+ # One bar per stage (white, filled green), all visible from the start
142
+ with Progress(
143
+ TextColumn("{task.description}"),
144
+ BarColumn(style="white", complete_style="green", finished_style="green"),
145
+ TaskProgressColumn(),
146
+ MofNCompleteColumn(),
147
+ ) as progress:
148
+ stage_tasks = {
149
+ split_name: progress.add_task(
150
+ f"[{stage}/{len(_SPLIT_NAMES)}] {split_name:<5}", total=len(planned_files[split_name])
151
+ )
152
+ for stage, split_name in enumerate(_SPLIT_NAMES, start=1)
153
+ }
154
+ for split_name in _SPLIT_NAMES:
155
+ new_labels_path = os.path.join(new_dataset_path, split_name, "labels")
156
+ new_images_path = os.path.join(new_dataset_path, split_name, "images")
157
+ for label_file, images_by_stem, old_to_new_index_match in planned_files[split_name]:
158
+ if not self.__copy_label_and_images(
159
+ label_file,
160
+ images_by_stem,
161
+ old_to_new_index_match,
162
+ new_labels_path,
163
+ new_images_path,
164
+ ignore_img_formats,
165
+ ):
166
+ return False
167
+ progress.advance(stage_tasks[split_name])
168
+
169
+ train_split, val_split, test_split = _SPLIT_NAMES
170
+ data_file_writer = YoloDataFileWriter(
171
+ new_dataset_path,
172
+ list(new_class_indices.keys()),
173
+ train_split=train_split,
174
+ val_split=val_split,
175
+ test_split=test_split,
176
+ )
177
+ try:
178
+ data_file_writer.write()
179
+ except Exception as e:
180
+ print(f"ERROR: Could not write new data.yaml to {new_dataset_path}. Generation of {data_set_name} is incomplete!")
181
+ print(f"Exception was: {e}")
182
+ return False
183
+
184
+ print(f"SUCCESS: Dataset '{os.path.basename(new_dataset_path)}' created at: {new_dataset_path}")
185
+ return True
186
+
187
+ def __copy_label_and_images(
188
+ self,
189
+ label_file: LabelFile,
190
+ images_by_stem: dict[str, list[str]],
191
+ old_to_new_index_match: dict[str, str],
192
+ new_labels_path: str,
193
+ new_images_path: str,
194
+ ignore_img_formats: bool,
195
+ ) -> bool:
196
+ """Copy one label file with remapped indices plus its images. Returns False on error."""
197
+ new_label_file_path = label_file.copy_by_class_indices(old_to_new_index_match, new_labels_path)
198
+ if not new_label_file_path:
199
+ return True
200
+
201
+ file_name_no_ext = os.path.splitext(os.path.basename(label_file.file_path))[0]
202
+ for img in images_by_stem.get(file_name_no_ext, []):
203
+ _, ext = os.path.splitext(img)
204
+
205
+ if not ignore_img_formats and ext.lower() not in self.SUPPORTED_IMAGE_FORMATS:
206
+ print(
207
+ f"ERROR: Image format not supported by yolo: {img}\n"
208
+ f"Supported formats are: {self.SUPPORTED_IMAGE_FORMATS}"
209
+ )
210
+ return False
211
+
212
+ new_image_name, _ = os.path.splitext(os.path.basename(new_label_file_path))
213
+ destination = os.path.join(new_images_path, new_image_name + ext)
214
+ try:
215
+ shutil.copyfile(img, destination)
216
+ except OSError as e:
217
+ print(
218
+ f"ERROR: Image could not be copied:\n"
219
+ f"Source: {img}\nDestination: {destination}\n"
220
+ f"Exception was: {e}"
221
+ )
222
+ return False
223
+ return True
224
+
225
+ def __generate_new_class_indices(self, class_names: set[str]) -> dict[str, str]:
226
+ """Generate a fresh label index mapping for the selected class names."""
227
+ ret: dict[str, str] = {}
228
+ for i, name in enumerate(class_names):
229
+ ret[name] = str(i)
230
+ return ret
231
+
232
+ def __generate_label_path_from_img_path(self, img_path: str) -> str:
233
+ """Convert an images directory path into its corresponding labels directory path."""
234
+ img_path = _normalize_path(img_path)
235
+ return os.path.normpath(os.path.join(os.path.dirname(img_path), "labels"))