LoadImages() pathlib update (#2140)

This commit is contained in:
Glenn Jocher 2021-02-05 09:06:23 -08:00 committed by GitHub
parent 73a0669930
commit 4bdc5a397e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 3 deletions

View File

@ -120,8 +120,7 @@ class _RepeatSampler(object):
class LoadImages: # for inference class LoadImages: # for inference
def __init__(self, path, img_size=640, stride=32): def __init__(self, path, img_size=640, stride=32):
p = str(Path(path)) # os-agnostic p = str(Path(path).absolute()) # os-agnostic absolute path
p = os.path.abspath(p) # absolute path
if '*' in p: if '*' in p:
files = sorted(glob.glob(p, recursive=True)) # glob files = sorted(glob.glob(p, recursive=True)) # glob
elif os.path.isdir(p): elif os.path.isdir(p):
@ -356,14 +355,17 @@ class LoadImagesAndLabels(Dataset): # for training/testing
p = Path(p) # os-agnostic p = Path(p) # os-agnostic
if p.is_dir(): # dir if p.is_dir(): # dir
f += glob.glob(str(p / '**' / '*.*'), recursive=True) f += glob.glob(str(p / '**' / '*.*'), recursive=True)
# f = list(p.rglob('**/*.*')) # pathlib
elif p.is_file(): # file elif p.is_file(): # file
with open(p, 'r') as t: with open(p, 'r') as t:
t = t.read().strip().splitlines() t = t.read().strip().splitlines()
parent = str(p.parent) + os.sep parent = str(p.parent) + os.sep
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)
else: else:
raise Exception(f'{prefix}{p} does not exist') raise Exception(f'{prefix}{p} does not exist')
self.img_files = sorted([x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in img_formats]) self.img_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
assert self.img_files, f'{prefix}No images found' assert self.img_files, f'{prefix}No images found'
except Exception as e: except Exception as e:
raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}') raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {help_url}')