Add YOLOv5 dataset yamls (#207)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Glenn Jocher 2023-01-10 18:33:44 +01:00 committed by GitHub
parent e371e81aa0
commit c7629e93bd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 2018 additions and 39 deletions

View file

@ -288,5 +288,6 @@ def check_dataset(dataset: str):
train_set = data_dir / "train"
test_set = data_dir / 'test' if (data_dir / 'test').exists() else data_dir / 'val' # data/test or data/val
nc = len([x for x in (data_dir / 'train').glob('*') if x.is_dir()]) # number of classes
names = [name for name in os.listdir(data_dir / 'train') if os.path.isdir(data_dir / 'train' / name)]
names = [x.name for x in (data_dir / 'train').iterdir() if x.is_dir()] # class names list
names = dict(enumerate(sorted(names)))
return {"train": train_set, "val": test_set, "nc": nc, "names": names}