Update datasets.py

This commit is contained in:
Glenn Jocher 2020-07-09 13:36:23 -07:00 committed by GitHub
parent 54a9e4f876
commit f9bab6b127
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 6 deletions

View File

@ -295,16 +295,11 @@ class LoadImagesAndLabels(Dataset): # for training/testing
else:
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]
path = subpath
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