Skip to content

Commit

Permalink
Merge pull request #4742 from imyhxy/fix-yolov5-importer
Browse files Browse the repository at this point in the history
Update `YOLOv5DatasetImporter` to support loading list of `.txt`
  • Loading branch information
brimoor authored Aug 28, 2024
2 parents e8ec301 + 14dfedf commit c9cfc05
Showing 1 changed file with 15 additions and 15 deletions.
30 changes: 15 additions & 15 deletions fiftyone/utils/yolo.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit c9cfc05

Please sign in to comment.