glenn-jocher commited on
Commit
c8a98cb
·
unverified ·
1 Parent(s): b1be685

Missing `nc` and `names` handling in check_dataset() (#4066)

Browse files
Files changed (1) hide show
  1. utils/general.py +3 -0
utils/general.py CHANGED
@@ -231,6 +231,9 @@ def check_dataset(data, autodownload=True):
231
  if data.get(k): # prepend path
232
  data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
233
 
 
 
 
234
  train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
235
  if val:
236
  val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
 
231
  if data.get(k): # prepend path
232
  data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]]
233
 
234
+ assert 'nc' in data, "Dataset 'nc' key missing."
235
+ if 'names' not in data:
236
+ data['names'] = [str(i) for i in range(data['nc'])] # assign class names if missing
237
  train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')]
238
  if val:
239
  val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path