Skip to content

Commit

Permalink
Improved dataset error introspection (ultralytics#8091)
Browse files Browse the repository at this point in the history
* Improved dataset error introspection

Intended to help ultralytics#8090

* Update general.py
  • Loading branch information
glenn-jocher authored and Clay Januhowski committed Sep 8, 2022
1 parent d5000df commit 3425577
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions utils/general.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ''
Expand All @@ -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():
Expand All @@ -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
Expand Down

0 comments on commit 3425577

Please sign in to comment.