zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Improved dataset error introspection (#8091) 

* Improved dataset error introspection

Intended to help https://github.com/ultralytics/yolov5/issues/8090

* Update general.py
This commit is contained in:
Glenn Jocher 2022-06-03 12:45:51 +02:00 committed by GitHub
parent 7cef03dddd
commit 47233e1698
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
utils/general.py
View File

@ -448,8 +448,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 = ''
@ -463,6 +462,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():
@ -472,9 +478,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