Browse Source

Refactor code to reduce duplication

5.0
lorenzomammana 4 years ago
parent
commit
54a9e4f876
1 changed files with 16 additions and 17 deletions
  1. +16
    -17
      utils/datasets.py

+ 16
- 17
utils/datasets.py View File

@@ -280,32 +280,31 @@ class LoadImagesAndLabels(Dataset): # for training/testing
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
cache_images=False, single_cls=False, stride=32, pad=0.0):
try:
if type(path) is list:
# Multiple datasets handler
f = []
for subpath in path:
f = []
for subpath in path if isinstance(path, list) else [path]:
subpath = str(Path(subpath)) # os-agnostic
parent = str(Path(subpath).parent) + os.sep
if os.path.isfile(subpath): # file
with open(subpath, 'r') as t:
subpath = str(Path(subpath)) # os-agnostic
parent = str(Path(subpath).parent) + os.sep
t = t.read().splitlines()
t = [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
f += t
path = str(Path(path[0])) # from now on treat multiple datasets as single
else:
path = str(Path(path)) # os-agnostic
parent = str(Path(path).parent) + os.sep
if os.path.isfile(path): # file
with open(path, 'r') as f:
f = f.read().splitlines()
f = [x.replace('./', parent) if x.startswith('./') else x for x in f] # local to global path
elif os.path.isdir(path): # folder
f = glob.iglob(path + os.sep + '*.*')
elif os.path.isdir(subpath): # folder
f = glob.iglob(subpath + os.sep + '*.*')
# Maybe change this to f += glob.glob, this should allow handling also multiple folders
else:
raise Exception('%s does not exist' % path)
raise Exception('%s does not exist' % subpath)
self.img_files = [x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats]
except:
# Maybe avoid handling bare exceptions
raise Exception('Error loading data from %s. See %s' % (path, help_url))

# Still need to do this for compatibility with the .npy and shape file saves
if isinstance(path, list):
path = str(Path(path[0]))
else:
path = str(Path(path))

n = len(self.img_files)
assert n > 0, 'No images found in %s. See %s' % (path, help_url)
bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index

Loading…
Cancel
Save