diff --git a/src/supervision/dataset/formats/yolo.py b/src/supervision/dataset/formats/yolo.py index 9abaff45..0e20bb7d 100644 --- a/src/supervision/dataset/formats/yolo.py +++ b/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( diff --git a/tests/dataset/formats/test_yolo.py b/tests/dataset/formats/test_yolo.py index 86047dc9..9a6e76ac 100644 --- a/tests/dataset/formats/test_yolo.py +++ b/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"), [