diff --git a/fiftyone/utils/yolo.py b/fiftyone/utils/yolo.py index 7f9840b144..db11ee015f 100644 --- a/fiftyone/utils/yolo.py +++ b/fiftyone/utils/yolo.py @@ -506,25 +506,25 @@ def setup(self): ) dataset_path = d.get("path", "") - data = fos.normpath(os.path.join(dataset_path, d[self.split])) + split_info = d[self.split] + if isinstance(split_info, str): + split_info = [split_info] + data = [ + fos.normpath(os.path.join(dataset_path, si)) for si in split_info + ] classes = _parse_yolo_classes(d.get("names", None)) - if etau.is_str(data) and data.endswith(".txt"): - txt_path = _parse_yolo_v5_path(data, self.yaml_path) - image_paths = [ - _parse_yolo_v5_path(fos.normpath(p), txt_path) - for p in _read_file_lines(txt_path) - ] - else: - if etau.is_str(data): - data_dirs = [data] + image_paths = [] + for data_path in data: + if etau.is_str(data_path) and data_path.endswith(".txt"): + txt_path = _parse_yolo_v5_path(data_path, self.yaml_path) + image_paths.extend( + _parse_yolo_v5_path(fos.normpath(p), txt_path) + for p in _read_file_lines(txt_path) + ) else: - data_dirs = data - - image_paths = [] - for data_dir in data_dirs: data_dir = fos.normpath( - _parse_yolo_v5_path(data_dir, self.yaml_path) + _parse_yolo_v5_path(data_path, self.yaml_path) ) image_paths.extend( etau.list_files(data_dir, abs_paths=True, recursive=True)