From 7849033da34bddbd0b231ad530d7f2387cb23fc4 Mon Sep 17 00:00:00 2001 From: hd Date: Tue, 27 Jun 2023 16:46:28 +0200 Subject: [PATCH] Fixed yaml file reading code --- supervision/dataset/formats/yolo.py | 18 +++--------------- supervision/utils/file.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/supervision/dataset/formats/yolo.py b/supervision/dataset/formats/yolo.py index 28dc4c59..67357613 100644 --- a/supervision/dataset/formats/yolo.py +++ b/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 diff --git a/supervision/utils/file.py b/supervision/utils/file.py index 7af669ca..14e5e014 100644 --- a/supervision/utils/file.py +++ b/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 \ No newline at end of file