Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update dataset_stats() for zipped datasets #3926

Merged
merged 2 commits into from
Jul 8, 2021
Merged

Update dataset_stats() for zipped datasets #3926

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a653576
Update `dataset_stats()` for zipped datasets
glenn-jocher Jul 7, 2021
a030269
cleanup
glenn-jocher Jul 7, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions utils/datasets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -888,9 +888,11 @@ def verify_image_label(args):

def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False):
""" Return dataset statistics dictionary with images and instances counts per split per class
Usage: from utils.datasets import *; dataset_stats('coco128.yaml', verbose=True)
Usage1: from utils.datasets import *; dataset_stats('coco128.yaml', verbose=True)
Usage2: from utils.datasets import *; dataset_stats('../datasets/coco128.zip', verbose=True)

Arguments
path: Path to data.yaml
path: Path to data.yaml or data.zip (with data.yaml inside data.zip)
autodownload: Attempt to download dataset if not found locally
verbose: Print stats dictionary
"""
Expand All @@ -899,8 +901,20 @@ def round_labels(labels):
# Update labels to integer class and 6 decimal place floats
return [[int(c), *[round(x, 6) for x in points]] for c, *points in labels]

with open(check_file(path)) as f:
def unzip(path):
# Unzip data.zip TODO: CONSTRAINT: path/to/abc.zip MUST unzip to 'path/to/abc/'
if str(path).endswith('.zip'): # path is data.zip
assert os.system(f'unzip -q {path} -d {path.parent}') == 0, f'Error unzipping {path}'
data_dir = path.with_suffix('') # dataset directory
return True, data_dir, list(data_dir.rglob('*.yaml'))[0] # zipped, data_dir, yaml_path
else: # path is data.yaml
return False, None, path

zipped, data_dir, yaml_path = unzip(Path(path))
with open(check_file(yaml_path)) as f:
data = yaml.safe_load(f) # data dict
if zipped:
data['path'] = data_dir # TODO: should this be dir.resolve()?
check_dataset(data, autodownload) # download dataset if missing
nc = data['nc'] # number of classes
stats = {'nc': nc, 'names': data['names']} # statistics dictionary
Expand Down