fix: populate class_name in DetectionDataset annotations (#2156)

DetectionDataset.__init__ now maps class_id to class names using
CLASS_NAME_DATA_FIELD so that LabelAnnotator displays human-readable
labels instead of raw integer IDs.

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>

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))
 

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
+                )