lorenzomammana commited on
Commit
54a9e4f
·
1 Parent(s): 94342ac

Refactor code to reduce duplication

Browse files
Files changed (1) hide show
  1. 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
- if type(path) is list:
284
- # Multiple datasets handler
285
- f = []
286
- for subpath in path:
 
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 = str(Path(path[0])) # from now on treat multiple datasets as single
294
- else:
295
- path = str(Path(path)) # os-agnostic
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' % path)
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