fix: sort YOLO class names by numeric keys (#2296)

* fix: sort yolo class names by numeric keys
* fix(yolo): reject double-hyphen keys in _is_int_like predicate
* fix(yolo): raise ValueError for mixed numeric/non-numeric names keys
* test(yolo): parametrize _extract_class_names with all key-type cases
* refine(yolo): rename lambda param key→k to avoid shadowing outer variable
* docs(yolo): comment bool guard in _is_int_like
* docs(yolo): docstring for _extract_class_names

---------

Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Claude Code <noreply@anthropic.com>

src/supervision/dataset/formats/yolo.py

@@ -68,6 +68,25 @@ def _with_seg_mask(lines: list[str]) -> bool:
 
 
 def _extract_class_names(file_path: str) -> list[str]:
+    """Return class names from a YOLO data.yaml file ordered by class index.
+
+    Supports list and dict forms of the ``names`` field. Dict keys that are
+    all int-like (plain ints or digit strings) are sorted numerically so
+    class index 10 follows index 9. All-non-numeric keys are sorted
+    lexicographically. Mixed numeric/non-numeric keys raise ``ValueError``.
+    Boolean YAML keys (``true``/``false``) are excluded from numeric sorting
+    because ``bool`` is a subclass of ``int`` in Python.
+
+    Args:
+        file_path: Path to the data.yaml file.
+
+    Returns:
+        Class names in class-index order.
+
+    Raises:
+        ValueError: If the YAML root is not a mapping, if ``names`` is
+            neither a list nor a dict, or if the dict has mixed key types.
+    """
     data: dict[str, Any] = read_yaml_file(file_path=file_path)
     if not isinstance(data, dict):
         raise ValueError(
@@ -76,7 +95,33 @@ def _extract_class_names(file_path: str) -> list[str]:
         )
     names = data.get("names")
     if isinstance(names, dict):
-        return [str(names[key]) for key in sorted(names.keys())]
+        keys = list(names.keys())
+
+        def _is_int_like(key: Any) -> bool:
+            # bool subclasses int; YAML `true`/`false` must not become class indices
+            if isinstance(key, bool):
+                return False
+            if isinstance(key, int):
+                return True
+            if isinstance(key, str):
+                stripped = key.strip()
+                return stripped.isdigit()
+            return False
+
+        int_like = [_is_int_like(k) for k in keys]
+        if any(int_like) and not all(int_like):
+            mixed_numeric = [k for k, il in zip(keys, int_like) if il][:3]
+            mixed_other = [k for k, il in zip(keys, int_like) if not il][:3]
+            raise ValueError(
+                f"Expected 'names' dict in data.yaml at '{file_path}' to have either "
+                f"all numeric or all non-numeric keys, got a mix: "
+                f"numeric {mixed_numeric} and non-numeric {mixed_other} keys."
+            )
+        if all(int_like):
+            sorted_keys = sorted(keys, key=lambda k: int(k))
+        else:
+            sorted_keys = sorted(keys, key=str)
+        return [str(names[key]) for key in sorted_keys]
     if isinstance(names, list):
         return [str(name) for name in names]
     raise ValueError(

tests/dataset/formats/test_yolo.py

@@ -10,6 +10,7 @@ import pytest
 from PIL import Image
 
 from supervision.dataset.formats.yolo import (
+    _extract_class_names,
     _image_name_to_annotation_name,
     _with_seg_mask,
     detections_to_yolo_annotations,
@@ -233,6 +234,50 @@ def test_image_name_to_annotation_name(
         assert result == expected_result
 
 
+@pytest.mark.parametrize(
+    ("yaml_text", "expected_names", "exception"),
+    [
+        (
+            "names:\n  '0': background\n  '1': person\n"
+            "  '2': car\n  '10': traffic_light\n",
+            ["background", "person", "car", "traffic_light"],
+            DoesNotRaise(),
+        ),  # quoted string numeric keys sort by integer value, not lexicographically
+        (
+            "names:\n  0: background\n  2: car\n  10: traffic_light\n",
+            ["background", "car", "traffic_light"],
+            DoesNotRaise(),
+        ),  # native int keys (most common YOLO format from Ultralytics/Roboflow)
+        (
+            "names:\n  cat: 0\n  dog: 1\n",
+            ["0", "1"],
+            DoesNotRaise(),
+        ),  # non-numeric string keys fall back to lexicographic sort
+        (
+            "names: {}\n",
+            [],
+            DoesNotRaise(),
+        ),  # empty names dict returns empty list
+        (
+            "names:\n  '--1': ignore\n  '0': person\n",
+            None,
+            pytest.raises(ValueError, match="mix"),
+        ),  # mixed numeric/non-numeric keys raise ValueError
+    ],
+)
+def test_extract_class_names_sorts_numeric_string_keys(
+    tmp_path: Path,
+    yaml_text: str,
+    expected_names: list[str] | None,
+    exception: Exception,
+) -> None:
+    """_extract_class_names returns class names sorted by class index."""
+    data_yaml_path = tmp_path / "data.yaml"
+    data_yaml_path.write_text(yaml_text, encoding="utf-8")
+    with exception:
+        assert _extract_class_names(file_path=str(data_yaml_path)) == expected_names
+
+
 @pytest.mark.parametrize(
     ("xyxy", "class_id", "image_shape", "polygon", "expected_result", "exception"),
     [