feat: make YOLOv5DatasetImporter support filelist source

voxel51 · Aug 28, 2024 · 475310d · 475310d
1 parent 9929c04
commit 475310d
Showing 1 changed file with 15 additions and 15 deletions.
diff --git 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)