mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Autosplit (#1488)

pull/1489/head
Glenn Jocher 2020-11-23 18:35:25 +01:00
parent 4798e66fdf
commit 354109c54c
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
utils/datasets.py
View File

@ -904,18 +904,19 @@ def flatten_recursive(path='../coco128'):
shutil.copyfile(file, new_path / Path(file).name) shutil.copyfile(file, new_path / Path(file).name)
def autosplit(path='../coco128', weights=(0.9, 0.1, 0.0)): # from utils.datasets import *; autosplit() def autosplit(path='../coco128', weights=(0.9, 0.1, 0.0)): # from utils.datasets import *; autosplit('../coco128')
""" Autosplit a dataset into train/val/test splits and save *.txt files """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
# Arguments # Arguments
path: Path to images directory path: Path to images directory
weights: Train, val, test weights (list) weights: Train, val, test weights (list)
""" """
path = Path(path) # images dir path = Path(path) # images dir
files = list(path.rglob('*.*')) files = list(path.rglob('*.*'))
indices = random.choices([0, 1, 2], weights=weights, k=len(files)) # assign each image to a split n = len(files) # number of files
indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split
txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files
[(path / x).unlink() for x in txt if (path / x).exists()] # remove existing [(path / x).unlink() for x in txt if (path / x).exists()] # remove existing
for i, img in tqdm(zip(indices, files)): for i, img in tqdm(zip(indices, files), total=n):
if img.suffix[1:] in img_formats: if img.suffix[1:] in img_formats:
with open(path / txt[i], 'a') as f: with open(path / txt[i], 'a') as f:
f.write(str(img) + '\n') # add image to txt file f.write(str(img) + '\n') # add image to txt file