Fixed yaml file reading code

supervision/dataset/formats/yolo.py

@@ -13,6 +13,7 @@ from supervision.utils.file import (
     list_files_with_extensions,
     read_txt_file,
     save_text_file,
+    read_yaml_file
 )
 
 
@@ -67,21 +68,8 @@ def _with_mask(lines: List[str]) -> bool:
 
 
 def _extract_class_names(file_path: str) -> List[str]:
-    names = []
-
-    with open(file_path, "r") as file:
-        lines = file.readlines()
-        for line in lines:
-            if line.strip().startswith("names:"):
-                start_index = lines.index(line) + 1
-                break
-
-        for line in lines[start_index:]:
-            if line.strip().startswith("-"):
-                names.append(line.strip().replace("-", "").strip())
-            else:
-                break
-
+    data = read_yaml_file(file_path=file_path)
+    names = data.get("names", [])
     return names
 
 

supervision/utils/file.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import List, Optional, Union
 
 import numpy as np
+import yaml
 
 
 class NumpyJsonEncoder(json.JSONEncoder):
@@ -107,3 +108,18 @@ def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
     """
     with open(file_path, "w") as fp:
         json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
+
+
+def read_yaml_file(file_path: str) -> dict:
+    """
+    Read a json file and return a dict.
+
+    Args:
+        file_path (str): The path to the json file.
+
+    Returns:
+        dict: A dict of annotations information
+    """
+    with open(file_path, "r") as file:
+        data = yaml.safe_load(file)
+    return data