lorenzomammana
commited on
Commit
·
54a9e4f
1
Parent(s):
94342ac
Refactor code to reduce duplication
Browse files- utils/datasets.py +16 -17
utils/datasets.py
CHANGED
@@ -280,32 +280,31 @@ class LoadImagesAndLabels(Dataset): # for training/testing
|
|
280 |
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
|
281 |
cache_images=False, single_cls=False, stride=32, pad=0.0):
|
282 |
try:
|
283 |
-
|
284 |
-
|
285 |
-
|
286 |
-
|
|
|
287 |
with open(subpath, 'r') as t:
|
288 |
-
subpath = str(Path(subpath)) # os-agnostic
|
289 |
-
parent = str(Path(subpath).parent) + os.sep
|
290 |
t = t.read().splitlines()
|
291 |
t = [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
292 |
f += t
|
293 |
-
path
|
294 |
-
|
295 |
-
|
296 |
-
parent = str(Path(path).parent) + os.sep
|
297 |
-
if os.path.isfile(path): # file
|
298 |
-
with open(path, 'r') as f:
|
299 |
-
f = f.read().splitlines()
|
300 |
-
f = [x.replace('./', parent) if x.startswith('./') else x for x in f] # local to global path
|
301 |
-
elif os.path.isdir(path): # folder
|
302 |
-
f = glob.iglob(path + os.sep + '*.*')
|
303 |
else:
|
304 |
-
raise Exception('%s does not exist' %
|
305 |
self.img_files = [x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats]
|
306 |
except:
|
|
|
307 |
raise Exception('Error loading data from %s. See %s' % (path, help_url))
|
308 |
|
|
|
|
|
|
|
|
|
|
|
|
|
309 |
n = len(self.img_files)
|
310 |
assert n > 0, 'No images found in %s. See %s' % (path, help_url)
|
311 |
bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
|
|
|
280 |
def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
|
281 |
cache_images=False, single_cls=False, stride=32, pad=0.0):
|
282 |
try:
|
283 |
+
f = []
|
284 |
+
for subpath in path if isinstance(path, list) else [path]:
|
285 |
+
subpath = str(Path(subpath)) # os-agnostic
|
286 |
+
parent = str(Path(subpath).parent) + os.sep
|
287 |
+
if os.path.isfile(subpath): # file
|
288 |
with open(subpath, 'r') as t:
|
|
|
|
|
289 |
t = t.read().splitlines()
|
290 |
t = [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path
|
291 |
f += t
|
292 |
+
elif os.path.isdir(subpath): # folder
|
293 |
+
f = glob.iglob(subpath + os.sep + '*.*')
|
294 |
+
# Maybe change this to f += glob.glob, this should allow handling also multiple folders
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
295 |
else:
|
296 |
+
raise Exception('%s does not exist' % subpath)
|
297 |
self.img_files = [x.replace('/', os.sep) for x in f if os.path.splitext(x)[-1].lower() in img_formats]
|
298 |
except:
|
299 |
+
# Maybe avoid handling bare exceptions
|
300 |
raise Exception('Error loading data from %s. See %s' % (path, help_url))
|
301 |
|
302 |
+
# Still need to do this for compatibility with the .npy and shape file saves
|
303 |
+
if isinstance(path, list):
|
304 |
+
path = str(Path(path[0]))
|
305 |
+
else:
|
306 |
+
path = str(Path(path))
|
307 |
+
|
308 |
n = len(self.img_files)
|
309 |
assert n > 0, 'No images found in %s. See %s' % (path, help_url)
|
310 |
bi = np.floor(np.arange(n) / batch_size).astype(np.int) # batch index
|