YoloDatasetCustomizer 0.3.0__tar.gz → 0.4.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: YoloDatasetCustomizer
-Version: 0.3.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
 
@@ -65,11 +72,11 @@ The project uses uv. Install dependencies and run tests with:
 
 # 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.3.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
 
@@ -38,11 +44,11 @@ The project uses uv. Install dependencies and run tests with:
 
 # 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.3.0 → yolodatasetcustomizer-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "YoloDatasetCustomizer"
-version = "0.3.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]

{yolodatasetcustomizer-0.3.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.3.0"
+__version__ = "0.4.0"
 __author__ = "Luca Mazzon"
 __email__ = "mazzon.luca@gmail.com"
 

{yolodatasetcustomizer-0.3.0 → yolodatasetcustomizer-0.4.0}/src/YoloDatasetCustomizer/customizer.py

@@ -3,6 +3,7 @@ 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
@@ -86,6 +87,16 @@ class YoloDatasetCustomizer:
         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:
@@ -112,46 +123,46 @@ class YoloDatasetCustomizer:
                         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)
+                    # 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)
-                        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()))
+                        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:
@@ -161,6 +172,44 @@ class YoloDatasetCustomizer:
 
         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] = {}

{yolodatasetcustomizer-0.3.0 → yolodatasetcustomizer-0.4.0}/test/YoloDatasetCustomizer_test.py

@@ -340,6 +340,27 @@ nc: 2
         new_image_path = os.path.join(new_dataset_path, "train", "images", "person_car.jpg")
         self.assertTrue(os.path.exists(new_image_path))
 
+    def test_merge_keeps_splits_separate_and_yaml_matches_folders(self):
+        customizer = YoloDatasetCustomizer([self.good_dataset_dir, self.reordered_names_dataset_dir])
+        name = "merged_splits_dataset"
+        self.assertTrue(
+            customizer.create_new_dataset_for_class_names({"person", "car"}, dst_path=self.workspace_root, data_set_name=name)
+        )
+        new_dataset_path = os.path.join(self.workspace_root, name)
+        with open(os.path.join(new_dataset_path, "data.yaml"), "r", encoding="utf-8") as f:
+            new_data = yaml.safe_load(f)
+
+        for key in ("train", "val", "test"):
+            self.assertTrue(os.path.isdir(os.path.join(new_dataset_path, new_data[key])), f"{key} -> {new_data[key]}")
+
+        # good_dataset: person_car in train, car in valid, person in test
+        def stems(split: str) -> set[str]:
+            return {os.path.splitext(n)[0].split("_")[0] for n in os.listdir(os.path.join(new_dataset_path, split, "labels"))}
+
+        self.assertIn("person", stems("test"))
+        self.assertNotIn("car", stems("test"))
+        self.assertNotIn("person", stems("val"))
+
     def test_yolo_dataset_customizer_returns_false_for_empty_class_names(self):
         customizer = YoloDatasetCustomizer([self.good_dataset_dir])
         self.assertFalse(customizer.create_new_dataset_for_class_names(set(), dst_path=self.workspace_root, data_set_name="empty_selection"))

{yolodatasetcustomizer-0.3.0 → yolodatasetcustomizer-0.4.0}/uv.lock

@@ -372,7 +372,7 @@ name = "exceptiongroup"
 version = "1.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+    { name = "typing-extensions" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
 wheels = [
@@ -405,7 +405,7 @@ name = "importlib-metadata"
 version = "9.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "zipp", marker = "python_full_version < '3.15'" },
+    { name = "zipp" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" }
 wheels = [
@@ -1078,10 +1078,11 @@ wheels = [
 
 [[package]]
 name = "yolodatasetcustomizer"
-version = "0.3.0"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "pyyaml" },
+    { name = "rich" },
 ]
 
 [package.dev-dependencies]
@@ -1093,7 +1094,10 @@ dev = [
 ]
 
 [package.metadata]
-requires-dist = [{ name = "pyyaml", specifier = ">=6.0.3" }]
+requires-dist = [
+    { name = "pyyaml", specifier = ">=6.0.3" },
+    { name = "rich", specifier = ">=13.0.0" },
+]
 
 [package.metadata.requires-dev]
 dev = [