YoloDatasetCustomizer 0.1.0__tar.gz → 0.2.0__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.
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/.gitattributes +0 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/.gitignore +0 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/.python-version +0 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/LICENSE +0 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/PKG-INFO +1 -1
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/README.md +0 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/pyproject.toml +2 -1
- yolodatasetcustomizer-0.2.0/resources/Improvements.txt +33 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/__init__.py +17 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/customizer.py +166 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/labels.py +71 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/paths.py +20 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/reader.py +90 -0
- yolodatasetcustomizer-0.2.0/src/YoloDatasetCustomizer/writer.py +41 -0
- {yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/test/YoloDatasetCustomizer_test.py +16 -1
- yolodatasetcustomizer-0.2.0/uv.lock +808 -0
- yolodatasetcustomizer-0.1.0/src/YoloDatasetCustomizer/YoloDatasetCustomizer.py +0 -351
- yolodatasetcustomizer-0.1.0/src/YoloDatasetCustomizer/__init__.py +0 -9
- yolodatasetcustomizer-0.1.0/uv.lock +0 -224
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: YoloDatasetCustomizer
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
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
|
|
File without changes
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "YoloDatasetCustomizer"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.2.0"
|
|
8
8
|
description = "Load one or multiple datasets. Reduce, mix or merge into a custom dataset."
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10"
|
|
@@ -40,4 +40,5 @@ Repository = "https://github.com/mLuca/YoloDatasetCustomizer"
|
|
|
40
40
|
[dependency-groups]
|
|
41
41
|
dev = [
|
|
42
42
|
"pytest>=9.1.1",
|
|
43
|
+
"twine>=6.2.0",
|
|
43
44
|
]
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
Priority-ordered improvement list
|
|
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
|
+
3. Tighten validation and error handling
|
|
10
|
+
Why: The project often prints warnings and continues instead of surfacing a clear, actionable failure. For dataset processing, silently skipping broken YAML or missing directories can produce partially valid outputs and make debugging much harder.
|
|
11
|
+
|
|
12
|
+
FIXED 4. Fix broken internal state and debug output
|
|
13
|
+
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.
|
|
14
|
+
|
|
15
|
+
5. Standardize class-index typing and conversion
|
|
16
|
+
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.
|
|
17
|
+
|
|
18
|
+
6. Refactor the monolithic implementation into smaller modules
|
|
19
|
+
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.
|
|
20
|
+
|
|
21
|
+
7. Improve test coverage around edge cases
|
|
22
|
+
Why: The current tests cover the happy path well, but not enough realistic failure cases such as malformed YAML, duplicate file names, empty split folders, and unexpected dataset layouts. More regression tests would reduce the chance of silent breakage.
|
|
23
|
+
|
|
24
|
+
8. Clarify and formally validate documented limitations
|
|
25
|
+
Why: The README already explains assumptions like fixed dataset naming and ignored YAML fields. Those constraints should be treated as explicit API requirements so users understand exactly what is supported and what is not.
|
|
26
|
+
|
|
27
|
+
9. Clean up naming and spelling inconsistencies
|
|
28
|
+
Why: Terms like “indeces” and inconsistent variable naming make the code less readable. This is not just cosmetic; it slows down maintenance and increases the chance of mistakes during refactors.
|
|
29
|
+
|
|
30
|
+
10. Add clearer public API contracts and examples
|
|
31
|
+
Why: The project is useful, but the usage patterns are not fully anchored in a stable interface. Clearer public contracts and examples would make the package easier to adopt and less fragile over time.
|
|
32
|
+
|
|
33
|
+
Highest priority overall: fix robustness around paths, defaults, and validation. Those are the issues most likely to lead to real-world failures in a dataset utility.
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""YOLO Dataset Customizer - Load, mix, merge, and customize YOLO datasets."""
|
|
2
|
+
|
|
3
|
+
from .customizer import YoloDatasetCustomizer
|
|
4
|
+
from .labels import LabelFile
|
|
5
|
+
from .reader import YoloDataFileReader
|
|
6
|
+
from .writer import YoloDataFileWriter
|
|
7
|
+
|
|
8
|
+
__version__ = "0.2.0"
|
|
9
|
+
__author__ = "Luca Mazzon"
|
|
10
|
+
__email__ = "mazzon.luca@gmail.com"
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"YoloDatasetCustomizer",
|
|
14
|
+
"YoloDataFileWriter",
|
|
15
|
+
"YoloDataFileReader",
|
|
16
|
+
"LabelFile",
|
|
17
|
+
]
|
|
@@ -0,0 +1,166 @@
|
|
|
1
|
+
import glob
|
|
2
|
+
import os
|
|
3
|
+
import shutil
|
|
4
|
+
|
|
5
|
+
import yaml
|
|
6
|
+
|
|
7
|
+
from .labels import LabelFile
|
|
8
|
+
from .paths import _SPLIT_NAMES, _get_unique_path, _normalize_path
|
|
9
|
+
from .reader import YoloDataFileReader
|
|
10
|
+
from .writer import YoloDataFileWriter
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class YoloDatasetCustomizer:
|
|
14
|
+
"""Find and customize YOLO datasets across one or more directory paths."""
|
|
15
|
+
|
|
16
|
+
def __init__(self, data_set_paths: list[str]):
|
|
17
|
+
"""Initialize with YAML dataset paths or directories to scan for YOLO datasets."""
|
|
18
|
+
self.__found_data_file_paths: set[str] = set()
|
|
19
|
+
self.__data_sets: list[YoloDataFileReader] = []
|
|
20
|
+
|
|
21
|
+
self.add_data_sets(data_set_paths)
|
|
22
|
+
self.SUPPORTED_IMAGE_FORMATS = [
|
|
23
|
+
".avif",
|
|
24
|
+
".bmp",
|
|
25
|
+
".dng",
|
|
26
|
+
".heic",
|
|
27
|
+
".jp2",
|
|
28
|
+
".jpeg",
|
|
29
|
+
".jpg",
|
|
30
|
+
".mpo",
|
|
31
|
+
".png",
|
|
32
|
+
".tif",
|
|
33
|
+
".tiff",
|
|
34
|
+
".webp",
|
|
35
|
+
]
|
|
36
|
+
|
|
37
|
+
def add_data_sets(self, data_set_paths: list[str]):
|
|
38
|
+
"""Discover YOLO data.yaml files from explicit paths or directory trees."""
|
|
39
|
+
data_set_paths = [_normalize_path(path) for path in data_set_paths]
|
|
40
|
+
|
|
41
|
+
self.__found_data_file_paths.update(
|
|
42
|
+
[path for path in data_set_paths if os.path.basename(path).lower() == "data.yaml"]
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
paths_without_file_ending = [
|
|
46
|
+
path for path in data_set_paths if os.path.basename(path).lower() != "data.yaml"
|
|
47
|
+
]
|
|
48
|
+
for path in paths_without_file_ending:
|
|
49
|
+
self.__found_data_file_paths.update(glob.glob(os.path.join(path, "**", "data.yaml"), recursive=True))
|
|
50
|
+
|
|
51
|
+
for path in self.__found_data_file_paths:
|
|
52
|
+
try:
|
|
53
|
+
ydfr = YoloDataFileReader(path)
|
|
54
|
+
self.__data_sets.append(ydfr)
|
|
55
|
+
except (yaml.YAMLError, TypeError, ValueError) as e:
|
|
56
|
+
print(f"ERROR: Issue loading YAML file at {path}: {e}")
|
|
57
|
+
except FileNotFoundError:
|
|
58
|
+
print(f"ERROR: Yaml file does not exist at {path}")
|
|
59
|
+
|
|
60
|
+
def get_found_data_file_paths(self) -> set[str]:
|
|
61
|
+
"""Return all discovered YOLO dataset YAML file paths."""
|
|
62
|
+
return self.__found_data_file_paths
|
|
63
|
+
|
|
64
|
+
def get_found_class_names(self) -> set[str]:
|
|
65
|
+
"""Return the union of all class names from discovered datasets."""
|
|
66
|
+
found_classes: set[str] = set()
|
|
67
|
+
for ydf in self.__data_sets:
|
|
68
|
+
found_classes.update(ydf.get_class_names())
|
|
69
|
+
return found_classes
|
|
70
|
+
|
|
71
|
+
def create_new_dataset_for_class_names(
|
|
72
|
+
self,
|
|
73
|
+
class_names: set[str],
|
|
74
|
+
dst_path: str = ".",
|
|
75
|
+
data_set_name: str = "new_dataset",
|
|
76
|
+
ignore_img_formats: bool = False,
|
|
77
|
+
) -> bool:
|
|
78
|
+
"""Create a filtered dataset containing only the requested classes."""
|
|
79
|
+
if len(class_names) == 0:
|
|
80
|
+
print("ERROR: No classes selected for new dataset.")
|
|
81
|
+
return False
|
|
82
|
+
|
|
83
|
+
new_dataset_path = os.path.join(_normalize_path(dst_path), data_set_name)
|
|
84
|
+
new_dataset_path = _get_unique_path(new_dataset_path)
|
|
85
|
+
|
|
86
|
+
new_class_indeces = self.__generate_new_class_indeces(class_names)
|
|
87
|
+
print(f"INFO: New class indeces: {new_class_indeces}")
|
|
88
|
+
|
|
89
|
+
for data_set in self.__data_sets:
|
|
90
|
+
common_class_names = set(data_set.get_class_names()) & class_names
|
|
91
|
+
if not common_class_names:
|
|
92
|
+
continue
|
|
93
|
+
|
|
94
|
+
old_to_new_index_match: dict[str, str] = {}
|
|
95
|
+
for class_name in common_class_names:
|
|
96
|
+
old_index = str(data_set.get_class_names().index(class_name))
|
|
97
|
+
new_index = new_class_indeces[class_name]
|
|
98
|
+
old_to_new_index_match[old_index] = new_index
|
|
99
|
+
|
|
100
|
+
for split_name in _SPLIT_NAMES:
|
|
101
|
+
for split_path in data_set.get_split_paths_for_split_name(split_name):
|
|
102
|
+
old_images_path = os.path.join(_normalize_path(data_set.get_file_dir()), split_path)
|
|
103
|
+
old_labels_path = self.__generate_label_path_from_img_path(old_images_path)
|
|
104
|
+
|
|
105
|
+
for path in [old_labels_path, old_images_path]:
|
|
106
|
+
if not os.path.exists(path):
|
|
107
|
+
print(f"WARNING: Path '{path}' does not exist! Skipping it.")
|
|
108
|
+
continue
|
|
109
|
+
|
|
110
|
+
new_labels_path = os.path.join(new_dataset_path, split_name, "labels")
|
|
111
|
+
new_images_path = os.path.join(new_dataset_path, split_name, "images")
|
|
112
|
+
os.makedirs(new_labels_path, exist_ok=True)
|
|
113
|
+
os.makedirs(new_images_path, exist_ok=True)
|
|
114
|
+
|
|
115
|
+
for original_label_file_path in glob.glob(os.path.join(old_labels_path, "*.txt")):
|
|
116
|
+
label_file = LabelFile(original_label_file_path)
|
|
117
|
+
new_label_file_path = label_file.copy_by_class_indeces(old_to_new_index_match, new_labels_path)
|
|
118
|
+
if new_label_file_path:
|
|
119
|
+
file_name = os.path.basename(original_label_file_path)
|
|
120
|
+
file_name_no_ext = file_name[: file_name.rfind(".")]
|
|
121
|
+
corresponding_image_files = glob.glob(os.path.join(old_images_path, f"{file_name_no_ext}.*"))
|
|
122
|
+
for img in corresponding_image_files:
|
|
123
|
+
try:
|
|
124
|
+
_, ext = os.path.splitext(img)
|
|
125
|
+
except Exception:
|
|
126
|
+
print(f"ERROR: Not a valid image, extension missing for: {img}")
|
|
127
|
+
return False
|
|
128
|
+
|
|
129
|
+
if not ignore_img_formats and ext.lower() not in self.SUPPORTED_IMAGE_FORMATS:
|
|
130
|
+
print(
|
|
131
|
+
f"ERROR: Image format not supported by yolo: {img}\n"
|
|
132
|
+
f"Supported formats are: {self.SUPPORTED_IMAGE_FORMATS}"
|
|
133
|
+
)
|
|
134
|
+
return False
|
|
135
|
+
|
|
136
|
+
new_image_name, _ = os.path.splitext(os.path.basename(new_label_file_path))
|
|
137
|
+
new_image_name = new_image_name + ext
|
|
138
|
+
destination = os.path.join(new_images_path, new_image_name)
|
|
139
|
+
if not shutil.copyfile(img, destination):
|
|
140
|
+
print(
|
|
141
|
+
f"ERROR: Image could not be copied:\n"
|
|
142
|
+
f"Source: {img}\nDestination: {destination}"
|
|
143
|
+
)
|
|
144
|
+
return False
|
|
145
|
+
|
|
146
|
+
data_file_writer = YoloDataFileWriter(new_dataset_path, list(new_class_indeces.keys()))
|
|
147
|
+
try:
|
|
148
|
+
data_file_writer.write()
|
|
149
|
+
except Exception as e:
|
|
150
|
+
print(f"ERROR: Could not write new data.yaml to {new_dataset_path}. Generation of {data_set_name} is incomplete!")
|
|
151
|
+
print(f"Exception was: {e}")
|
|
152
|
+
return False
|
|
153
|
+
|
|
154
|
+
return True
|
|
155
|
+
|
|
156
|
+
def __generate_new_class_indeces(self, class_names: set[str]) -> dict[str, str]:
|
|
157
|
+
"""Generate a fresh label index mapping for the selected class names."""
|
|
158
|
+
ret: dict[str, str] = {}
|
|
159
|
+
for i, name in enumerate(class_names):
|
|
160
|
+
ret[name] = str(i)
|
|
161
|
+
return ret
|
|
162
|
+
|
|
163
|
+
def __generate_label_path_from_img_path(self, img_path: str) -> str:
|
|
164
|
+
"""Convert an images directory path into its corresponding labels directory path."""
|
|
165
|
+
img_path = _normalize_path(img_path)
|
|
166
|
+
return os.path.normpath(os.path.join(os.path.dirname(img_path), "labels"))
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import re
|
|
3
|
+
|
|
4
|
+
from .paths import _get_unique_path
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class LabelFile:
|
|
8
|
+
"""Read YOLO label files and copy label lines with remapped class indices."""
|
|
9
|
+
|
|
10
|
+
def __init__(self, file_path: str):
|
|
11
|
+
"""
|
|
12
|
+
Args:
|
|
13
|
+
file_path: Valid path to the label file
|
|
14
|
+
Raises:
|
|
15
|
+
FileNotFoundError: When the label file does not exist
|
|
16
|
+
"""
|
|
17
|
+
self.__class_indeces: set[str] = set()
|
|
18
|
+
self.file_path = _get_unique_path(file_path) if os.path.isdir(file_path) else os.path.abspath(file_path)
|
|
19
|
+
if not os.path.exists(self.file_path) or not os.path.isfile(self.file_path):
|
|
20
|
+
raise FileNotFoundError(f"No label file found at: {self.file_path}")
|
|
21
|
+
|
|
22
|
+
self.read_label_file()
|
|
23
|
+
|
|
24
|
+
def read_label_file(self):
|
|
25
|
+
"""Parse the label file and collect all class indices referenced in it."""
|
|
26
|
+
with open(self.file_path, encoding="utf-8") as f:
|
|
27
|
+
for i, line in enumerate(f):
|
|
28
|
+
match = re.match(r"^(\d+) ", line)
|
|
29
|
+
if not match:
|
|
30
|
+
print(f"WARNING: No class found on line {i} in label file {self.file_path}")
|
|
31
|
+
else:
|
|
32
|
+
self.__class_indeces.add(match.group(1))
|
|
33
|
+
|
|
34
|
+
def get_class_indices(self) -> set[str]:
|
|
35
|
+
"""Return the set of class indices in the loaded label file."""
|
|
36
|
+
return self.__class_indeces
|
|
37
|
+
|
|
38
|
+
def copy_by_class_indeces(self, old_to_new_index_match: dict[str, str], dst_path: str, file_name: str = "") -> str | None:
|
|
39
|
+
"""Copy the label file to the destination with indices remapped for the new dataset."""
|
|
40
|
+
common_class_indeces = self.__class_indeces & set(old_to_new_index_match.keys())
|
|
41
|
+
if not common_class_indeces:
|
|
42
|
+
return None
|
|
43
|
+
|
|
44
|
+
dst_path = os.path.abspath(dst_path)
|
|
45
|
+
if os.path.basename(dst_path) != "labels":
|
|
46
|
+
dst_path = os.path.join(dst_path, "labels")
|
|
47
|
+
|
|
48
|
+
if os.path.exists(dst_path) and os.path.isdir(dst_path):
|
|
49
|
+
new_content: list[str] = []
|
|
50
|
+
with open(self.file_path, encoding="utf-8") as f:
|
|
51
|
+
for line in f:
|
|
52
|
+
match = re.match(r"^(\d+) ", line)
|
|
53
|
+
if not match:
|
|
54
|
+
continue
|
|
55
|
+
elif match.group(1) in common_class_indeces:
|
|
56
|
+
old_index = match.group(1)
|
|
57
|
+
new_index = old_to_new_index_match[old_index]
|
|
58
|
+
if new_index != old_index:
|
|
59
|
+
line = line.replace(old_index, new_index, 1)
|
|
60
|
+
new_content.append(line)
|
|
61
|
+
|
|
62
|
+
if file_name == "":
|
|
63
|
+
file_name = os.path.basename(self.file_path)
|
|
64
|
+
dst_file_path = _get_unique_path(os.path.join(dst_path, file_name))
|
|
65
|
+
|
|
66
|
+
with open(dst_file_path, "w", encoding="utf-8") as f:
|
|
67
|
+
f.writelines(new_content)
|
|
68
|
+
else:
|
|
69
|
+
raise FileNotFoundError(f"Directory not found: {dst_path}")
|
|
70
|
+
|
|
71
|
+
return dst_file_path
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
import os
|
|
2
|
+
|
|
3
|
+
_SPLIT_NAMES = ["train", "val", "test"]
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def _normalize_path(path: str) -> str:
|
|
7
|
+
"""Normalize path separators and return an absolute path for cross-platform use."""
|
|
8
|
+
return os.path.abspath(str(path).replace("\\", os.sep))
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def _get_unique_path(path: str) -> str:
|
|
12
|
+
"""Return an unused file path by appending a numeric suffix when needed."""
|
|
13
|
+
path = _normalize_path(path)
|
|
14
|
+
base, ext = os.path.splitext(path)
|
|
15
|
+
unique_path = path
|
|
16
|
+
index = 1
|
|
17
|
+
while os.path.exists(unique_path):
|
|
18
|
+
unique_path = f"{base}_{index}{ext}"
|
|
19
|
+
index += 1
|
|
20
|
+
return unique_path
|
|
@@ -0,0 +1,90 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from typing import Any
|
|
3
|
+
|
|
4
|
+
import yaml
|
|
5
|
+
|
|
6
|
+
from .paths import _SPLIT_NAMES, _normalize_path
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class YoloDataFileReader:
|
|
10
|
+
"""Read a YOLO dataset YAML file and expose its split paths and class names."""
|
|
11
|
+
|
|
12
|
+
def __init__(self, file_path: str):
|
|
13
|
+
"""
|
|
14
|
+
Args:
|
|
15
|
+
file_path: Valid path to the data set YAML file
|
|
16
|
+
Raises:
|
|
17
|
+
yaml.YAMLError: When there is an issue with reading the data-set YAML file
|
|
18
|
+
FileNotFoundError: When the data-set YAML file does not exist
|
|
19
|
+
ValueError: When a value in the YAML file is missing or is faulty
|
|
20
|
+
TypeError: When a value comes in an unexpected type
|
|
21
|
+
"""
|
|
22
|
+
self.__file_path = _normalize_path(file_path)
|
|
23
|
+
self.__file_dir = os.path.dirname(self.__file_path)
|
|
24
|
+
self.__splits_paths: dict[str, list[str]] = {}
|
|
25
|
+
self.__class_names: list[str] = []
|
|
26
|
+
|
|
27
|
+
self.sync_content()
|
|
28
|
+
|
|
29
|
+
def get_split_paths_for_split_name(self, split_name: str) -> list[str]:
|
|
30
|
+
"""Return a list of dataset split paths for the requested split name."""
|
|
31
|
+
return self.__splits_paths.get(split_name, [])
|
|
32
|
+
|
|
33
|
+
def get_class_names(self) -> list[str]:
|
|
34
|
+
"""Return the class names defined in the dataset YAML."""
|
|
35
|
+
return self.__class_names
|
|
36
|
+
|
|
37
|
+
def get_file_path(self) -> str:
|
|
38
|
+
"""Return the absolute path to the loaded YAML file."""
|
|
39
|
+
return self.__file_path
|
|
40
|
+
|
|
41
|
+
def get_file_dir(self) -> str:
|
|
42
|
+
"""Return the directory containing the loaded dataset YAML file."""
|
|
43
|
+
return self.__file_dir
|
|
44
|
+
|
|
45
|
+
def get_indices_for_names(self, class_names: set[str]) -> set[str]:
|
|
46
|
+
"""Map a set of class names to their corresponding numeric label indices."""
|
|
47
|
+
ret: set[str] = set()
|
|
48
|
+
for class_name in class_names:
|
|
49
|
+
try:
|
|
50
|
+
idx = self.__class_names.index(class_name)
|
|
51
|
+
ret.add(str(idx))
|
|
52
|
+
except ValueError:
|
|
53
|
+
pass
|
|
54
|
+
return ret
|
|
55
|
+
|
|
56
|
+
def sync_content(self):
|
|
57
|
+
"""Reload YAML content and refresh split and name mappings."""
|
|
58
|
+
with open(self.__file_path, encoding="utf-8") as f:
|
|
59
|
+
content = yaml.safe_load(f)
|
|
60
|
+
for split_name in _SPLIT_NAMES:
|
|
61
|
+
self.__splits_paths[split_name] = self.__read_yaml_entry(content, split_name)
|
|
62
|
+
self.__class_names = self.__read_yaml_entry(content, "names")
|
|
63
|
+
|
|
64
|
+
def __read_yaml_entry(self, yaml_content: Any, entry_name: str, missing_ok: bool = True) -> list[str]:
|
|
65
|
+
"""Normalize YAML entry values into a list of paths or names."""
|
|
66
|
+
entry = yaml_content.get(entry_name)
|
|
67
|
+
ret: list[str] = []
|
|
68
|
+
if entry:
|
|
69
|
+
if isinstance(entry, str):
|
|
70
|
+
ret.append(entry)
|
|
71
|
+
elif isinstance(entry, dict):
|
|
72
|
+
ret.extend(entry.values())
|
|
73
|
+
elif isinstance(entry, list):
|
|
74
|
+
ret.extend(entry)
|
|
75
|
+
else:
|
|
76
|
+
raise TypeError(f"Unexpected type for {entry_name} in: {self.__file_path}")
|
|
77
|
+
else:
|
|
78
|
+
if not missing_ok:
|
|
79
|
+
raise ValueError(f"{entry_name} not defined in: {self.__file_path}")
|
|
80
|
+
|
|
81
|
+
return ret
|
|
82
|
+
|
|
83
|
+
def __repr__(self) -> str:
|
|
84
|
+
return (
|
|
85
|
+
f"\nYoloDatasetFile(\n"
|
|
86
|
+
f" path: {self.__file_path}\n"
|
|
87
|
+
f" split_paths: {self.__splits_paths}\n"
|
|
88
|
+
f" class_names: {self.__class_names}\n"
|
|
89
|
+
f")\n"
|
|
90
|
+
)
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
import os
|
|
2
|
+
|
|
3
|
+
import yaml
|
|
4
|
+
|
|
5
|
+
from .paths import _normalize_path
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class YoloDataFileWriter:
|
|
9
|
+
"""Create and write a YOLO dataset YAML file for a generated dataset."""
|
|
10
|
+
|
|
11
|
+
def __init__(
|
|
12
|
+
self,
|
|
13
|
+
data_set_dir: str = "./",
|
|
14
|
+
class_names: list[str] | None = None,
|
|
15
|
+
train_split: str = "./train",
|
|
16
|
+
val_split: str = "./valid",
|
|
17
|
+
test_split: str = "./test",
|
|
18
|
+
) -> None:
|
|
19
|
+
"""Initialize writer with dataset directory, class names, and split locations."""
|
|
20
|
+
self.data_set_dir = _normalize_path(data_set_dir)
|
|
21
|
+
self.class_names = list(class_names) if class_names is not None else []
|
|
22
|
+
self.train_split = train_split
|
|
23
|
+
self.val_split = val_split
|
|
24
|
+
self.test_split = test_split
|
|
25
|
+
self.FILE_NAME = "data.yaml"
|
|
26
|
+
|
|
27
|
+
def write(self):
|
|
28
|
+
"""Write the YOLO data.yaml file with train/val/test paths and class names."""
|
|
29
|
+
classes = {}
|
|
30
|
+
for i, class_name in enumerate(self.class_names):
|
|
31
|
+
classes[i] = class_name
|
|
32
|
+
data = {
|
|
33
|
+
"train": f"{self.train_split}/images",
|
|
34
|
+
"val": f"{self.val_split}/images",
|
|
35
|
+
"test": f"{self.test_split}/images",
|
|
36
|
+
"names": classes,
|
|
37
|
+
"nc": len(self.class_names),
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
with open(os.path.join(self.data_set_dir, self.FILE_NAME), "w", encoding="utf-8") as f:
|
|
41
|
+
yaml.safe_dump(data, f, default_flow_style=False)
|
{yolodatasetcustomizer-0.1.0 → yolodatasetcustomizer-0.2.0}/test/YoloDatasetCustomizer_test.py
RENAMED
|
@@ -7,7 +7,7 @@ import yaml
|
|
|
7
7
|
|
|
8
8
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
|
|
9
9
|
|
|
10
|
-
from YoloDatasetCustomizer
|
|
10
|
+
from YoloDatasetCustomizer import (
|
|
11
11
|
YoloDataFileWriter,
|
|
12
12
|
YoloDataFileReader,
|
|
13
13
|
LabelFile,
|
|
@@ -130,6 +130,15 @@ nc: 1
|
|
|
130
130
|
self.assertEqual(loaded["names"], {0: "person", 1: "car"})
|
|
131
131
|
self.assertEqual(loaded["nc"], 2)
|
|
132
132
|
|
|
133
|
+
def test_yolo_data_file_writer_uses_independent_default_class_names(self):
|
|
134
|
+
writer_a = YoloDataFileWriter()
|
|
135
|
+
writer_b = YoloDataFileWriter()
|
|
136
|
+
|
|
137
|
+
writer_a.class_names.append("person")
|
|
138
|
+
|
|
139
|
+
self.assertEqual(writer_a.class_names, ["person"])
|
|
140
|
+
self.assertEqual(writer_b.class_names, [])
|
|
141
|
+
|
|
133
142
|
def test_yolo_data_file_reader_reads_splits_names_and_indices(self):
|
|
134
143
|
reader = YoloDataFileReader(os.path.join(self.good_dataset_dir, "data.yaml"))
|
|
135
144
|
self.assertEqual(reader.get_file_path(), os.path.abspath(os.path.join(self.good_dataset_dir, "data.yaml")))
|
|
@@ -184,6 +193,12 @@ nc: 1
|
|
|
184
193
|
self.assertIn(os.path.abspath(os.path.join(self.unsupported_image_dataset_dir, "data.yaml")), found_paths)
|
|
185
194
|
self.assertEqual(customizer.get_found_class_names(), {"person", "car"})
|
|
186
195
|
|
|
196
|
+
def test_yolo_dataset_customizer_handles_windows_style_paths(self):
|
|
197
|
+
windows_style_root = self.workspace_root.replace(os.sep, "\\")
|
|
198
|
+
customizer = YoloDatasetCustomizer([windows_style_root])
|
|
199
|
+
self.assertIn(os.path.abspath(os.path.join(self.good_dataset_dir, "data.yaml")), customizer.get_found_data_file_paths())
|
|
200
|
+
self.assertEqual(customizer.get_found_class_names(), {"person", "car"})
|
|
201
|
+
|
|
187
202
|
def test_yolo_dataset_customizer_creates_filtered_dataset(self):
|
|
188
203
|
customizer = YoloDatasetCustomizer([self.good_dataset_dir])
|
|
189
204
|
success = customizer.create_new_dataset_for_class_names({"car"}, dst_path=self.workspace_root, data_set_name="filtered_car_dataset")
|