diff --git a/utils/general.py b/utils/general.py index c33fe60d59f3..81ab327c7d53 100755 --- a/utils/general.py +++ b/utils/general.py @@ -449,8 +449,7 @@ def check_font(font=FONT, progress=False): def check_dataset(data, autodownload=True): - # Download and/or unzip dataset if not found locally - # Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip + # Download, check and/or unzip dataset if not found locally # Download (optional) extract_dir = '' @@ -464,6 +463,13 @@ def check_dataset(data, autodownload=True): with open(data, errors='ignore') as f: data = yaml.safe_load(f) # dictionary + # Checks + for k in 'train', 'val', 'nc': + assert k in data, emojis(f"data.yaml '{k}:' field missing ❌") + if 'names' not in data: + LOGGER.warning(emojis("data.yaml 'names:' field missing ⚠, assigning default names 'class0', 'class1', etc.")) + data['names'] = [f'class{i}' for i in range(data['nc'])] # default names + # Resolve paths path = Path(extract_dir or data.get('path') or '') # optional 'path' default to '.' if not path.is_absolute(): @@ -473,9 +479,6 @@ def check_dataset(data, autodownload=True): data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]] # Parse yaml - assert 'nc' in data, "Dataset 'nc' key missing." - if 'names' not in data: - data['names'] = [f'class{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