Unverified 提交 526e6505 authored 作者: Colin Wong's avatar Colin Wong 提交者: GitHub

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: 's avatarColin Wong <noreply@brains4drones.com> Co-authored-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com>
上级 c215e873
...@@ -176,15 +176,17 @@ class _RepeatSampler: ...@@ -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): ...@@ -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'
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论