From 413abc5749faeb415736607906a8385b594f3aa9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 19 Jul 2021 13:05:06 +0200 Subject: [PATCH] Missing `nc` and `names` handling in check_dataset() --- utils/general.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/utils/general.py b/utils/general.py index 846c1464c28c..08a3ff6539b2 100755 --- a/utils/general.py +++ b/utils/general.py @@ -231,6 +231,9 @@ def check_dataset(data, autodownload=True): if data.get(k): # prepend path data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]] + assert 'nc' in data, "Dataset 'nc' key missing." + if 'names' not in data: + data['names'] = [str(i) for i in range(data['nc'])] # assign class names if missing train, val, test, s = [data.get(x) for x in ('train', 'val', 'test', 'download')] if val: val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path