diff --git a/src/supervision/dataset/core.py b/src/supervision/dataset/core.py index 823339ab..44a98079 100644 --- a/src/supervision/dataset/core.py +++ b/src/supervision/dataset/core.py @@ -12,6 +12,7 @@ import numpy as np import numpy.typing as npt from supervision.classification.core import Classifications +from supervision.config import CLASS_NAME_DATA_FIELD from supervision.dataset.formats.coco import ( load_coco_annotations, save_coco_annotations, @@ -84,6 +85,14 @@ class DetectionDataset(BaseDataset): ) self.annotations = annotations + if self.classes: + np_classes = np.array(self.classes) + for annotation in self.annotations.values(): + if annotation.class_id is not None: + annotation.data[CLASS_NAME_DATA_FIELD] = np_classes[ + annotation.class_id + ] + # Eliminate duplicates while preserving order self.image_paths = list(dict.fromkeys(images)) diff --git a/tests/dataset/test_core.py b/tests/dataset/test_core.py index 451c3b0c..fcfa48e8 100644 --- a/tests/dataset/test_core.py +++ b/tests/dataset/test_core.py @@ -1,11 +1,14 @@ from __future__ import annotations from contextlib import ExitStack as DoesNotRaise +from pathlib import Path +import numpy as np import pytest -from supervision import DetectionDataset -from tests.helpers import _create_detections +from supervision import DetectionDataset, Detections +from supervision.config import CLASS_NAME_DATA_FIELD +from tests.helpers import _create_detections, create_yolo_dataset @pytest.mark.parametrize( @@ -187,3 +190,97 @@ def test_dataset_merge( with exception: result = DetectionDataset.merge(dataset_list=dataset_list) assert result == expected_result + + +class TestClassNamePopulation: + """Verify that DetectionDataset populates CLASS_NAME_DATA_FIELD on init.""" + + def test_class_name_populated_on_init(self) -> None: + """Basic case: class_name data field is set from classes and class_id.""" + dataset = DetectionDataset( + classes=["dog", "cat"], + images=["img1.png"], + annotations={ + "img1.png": _create_detections( + xyxy=[[0, 0, 10, 10], [20, 20, 30, 30]], + class_id=[0, 1], + ), + }, + ) + annotation = dataset.annotations["img1.png"] + assert CLASS_NAME_DATA_FIELD in annotation.data + np.testing.assert_array_equal( + annotation.data[CLASS_NAME_DATA_FIELD], + np.array(["dog", "cat"]), + ) + + def test_class_name_with_empty_annotations(self) -> None: + """Empty Detections should not raise an error.""" + dataset = DetectionDataset( + classes=["dog"], + images=["img1.png"], + annotations={"img1.png": Detections.empty()}, + ) + annotation = dataset.annotations["img1.png"] + assert CLASS_NAME_DATA_FIELD in annotation.data + assert len(annotation.data[CLASS_NAME_DATA_FIELD]) == 0 + + def test_class_name_with_empty_classes(self) -> None: + """When classes is empty, class_name should not be populated.""" + dataset = DetectionDataset( + classes=[], + images=[], + annotations={}, + ) + assert len(dataset.annotations) == 0 + + def test_class_name_after_merge(self) -> None: + """After merging datasets, class_name must match remapped class_id.""" + ds1 = DetectionDataset( + classes=["dog", "person"], + images=["img1.png"], + annotations={ + "img1.png": _create_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]), + }, + ) + ds2 = DetectionDataset( + classes=["cat"], + images=["img2.png"], + annotations={ + "img2.png": _create_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]), + }, + ) + merged = DetectionDataset.merge([ds1, ds2]) + + # merged.classes is ["cat", "dog", "person"] + # ds1's dog (0) -> dog (1), ds2's cat (0) -> cat (0) + ann1 = merged.annotations["img1.png"] + assert CLASS_NAME_DATA_FIELD in ann1.data + np.testing.assert_array_equal( + ann1.data[CLASS_NAME_DATA_FIELD], np.array(["dog"]) + ) + + ann2 = merged.annotations["img2.png"] + assert CLASS_NAME_DATA_FIELD in ann2.data + np.testing.assert_array_equal( + ann2.data[CLASS_NAME_DATA_FIELD], np.array(["cat"]) + ) + + def test_class_name_from_yolo(self, tmp_path: Path) -> None: + """Integration test: from_yolo should produce class_name data.""" + dataset_info = create_yolo_dataset( + str(tmp_path), num_images=2, classes=["cat", "dog"] + ) + dataset = DetectionDataset.from_yolo( + images_directory_path=dataset_info["images_dir"], + annotations_directory_path=dataset_info["labels_dir"], + data_yaml_path=dataset_info["data_yaml_path"], + ) + + for _, annotation in dataset.annotations.items(): + if annotation.class_id is not None and len(annotation.class_id) > 0: + assert CLASS_NAME_DATA_FIELD in annotation.data + expected_names = np.array(dataset.classes)[annotation.class_id] + np.testing.assert_array_equal( + annotation.data[CLASS_NAME_DATA_FIELD], expected_names + )