Fix LoadImages() with dataset YAML lists (#8517)
* Fix LoadImages with dataset yaml lists * Update dataloaders.py * Update dataloaders.py * Simplify/refactor PR * Update dataloaders.py Co-authored-by: Colin Wong <noreply@brains4drones.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
1 file changed
+12
-10
+12
-10
@@ -176,15 +176,17 @@ class _RepeatSampler:
|
|||||||
class LoadImages:
|
class LoadImages:
|
||||||
# YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
|
# YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4`
|
||||||
def __init__(self, path, img_size=640, stride=32, auto=True):
|
def __init__(self, path, img_size=640, stride=32, auto=True):
|
||||||
p = str(Path(path).resolve()) # os-agnostic absolute path
|
files = []
|
||||||
if '*' in p:
|
for p in sorted(path) if isinstance(path, (list, tuple)) else [path]:
|
||||||
files = sorted(glob.glob(p, recursive=True)) # glob
|
p = str(Path(p).resolve())
|
||||||
elif os.path.isdir(p):
|
if '*' in p:
|
||||||
files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir
|
files.extend(sorted(glob.glob(p, recursive=True))) # glob
|
||||||
elif os.path.isfile(p):
|
elif os.path.isdir(p):
|
||||||
files = [p] # files
|
files.extend(sorted(glob.glob(os.path.join(p, '*.*')))) # dir
|
||||||
else:
|
elif os.path.isfile(p):
|
||||||
raise Exception(f'ERROR: {p} does not exist')
|
files.append(p) # files
|
||||||
|
else:
|
||||||
|
raise FileNotFoundError(f'{p} does not exist')
|
||||||
|
|
||||||
images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
|
images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS]
|
||||||
videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
|
videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS]
|
||||||
@@ -437,7 +439,7 @@ class LoadImagesAndLabels(Dataset):
|
|||||||
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
||||||
# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
|
# f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib)
|
||||||
else:
|
else:
|
||||||
raise Exception(f'{prefix}{p} does not exist')
|
raise FileNotFoundError(f'{prefix}{p} does not exist')
|
||||||
self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
|
self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
|
||||||
# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
|
# self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib
|
||||||
assert self.im_files, f'{prefix}No images found'
|
assert self.im_files, f'{prefix}No images found'
|
||||||
|
|||||||
Reference in new issue
Block a user