democat Colin Wong glenn-jocher commited on
Commit
526e650
1 Parent(s): c215e87

Fix `LoadImages()` with dataset YAML lists (#8517)

Browse files

* Fix LoadImages with dataset yaml lists

* Update dataloaders.py

* Update dataloaders.py

* Simplify/refactor PR

* Update dataloaders.py

Co-authored-by: Colin Wong <[email protected]>
Co-authored-by: Glenn Jocher <[email protected]>

Files changed (1) hide show
  1. utils/dataloaders.py +12 -10
utils/dataloaders.py CHANGED
@@ -176,15 +176,17 @@ class _RepeatSampler:
176
  class LoadImages:
177
  # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
178
  def __init__(self, path, img_size=640, stride=32, auto=True):
179
- p = str(Path(path).resolve()) # os-agnostic absolute path
180
- if '*' in p:
181
- files = sorted(glob.glob(p, recursive=True)) # glob
182
- elif os.path.isdir(p):
183
- files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
184
- elif os.path.isfile(p):
185
- files = [p] # files
186
- else:
187
- raise Exception(f'ERROR: {p} does not exist')
 
 
188
 
189
  images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
190
  videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
@@ -437,7 +439,7 @@ class LoadImagesAndLabels(Dataset):
437
  f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
438
  # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
439
  else:
440
- raise Exception(f'{prefix}{p} does not exist')
441
  self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
442
  # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
443
  assert self.im_files, f'{prefix}No images found'
 
176
  class LoadImages:
177
  # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
178
  def __init__(self, path, img_size=640, stride=32, auto=True):
179
+ files = []
180
+ for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
181
+ p = str(Path(p).resolve())
182
+ if '*' in p:
183
+ files.extend(sorted(glob.glob(p, recursive=True))) # glob
184
+ elif os.path.isdir(p):
185
+ files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) # dir
186
+ elif os.path.isfile(p):
187
+ files.append(p) # files
188
+ else:
189
+ raise FileNotFoundError(f'{p} does not exist')
190
 
191
  images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
192
  videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
 
439
  f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
440
  # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
441
  else:
442
+ raise FileNotFoundError(f'{prefix}{p} does not exist')
443
  self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
444
  # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
445
  assert self.im_files, f'{prefix}No images found'