Spaces:
Runtime error
Runtime error
File size: 1,354 Bytes
1a24a58 cd1bbe4 1f73e71 cd1bbe4 3d78160 6d7f6ee 3d78160 6d7f6ee 3d78160 6d7f6ee e04a33d cd1bbe4 6d7f6ee 94a3eb0 3d78160 cd1bbe4 acbdf2a cd1bbe4 3d78160 1f73e71 bbaa3e0 6d7f6ee 26ef429 bbaa3e0 1f73e71 3d78160 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import os
import csv
import entity
import argparse
from common import defaults
YOLO_DATA_TEMPLATES = {
5: '''
train: %%datapath%%/split/train
val: %%datapath%%/split/val
''',
6: '''
train: %%datapath%%/split/images/train
val: %%datapath%%/split/images/val
test: %%datapath%%/split/images/test
is_coco: False
'''
}
def gen_data_yaml(bcos, datapath='../data', version=6):
names = [f"{d.name}" for d in bcos.values()]
return f'''
# this file is autogenerated by write_data.py for YOLO version {version}
{YOLO_DATA_TEMPLATES[version].replace('%%datapath%%', datapath)}
nc: {len(bcos.keys())}
names: {names}
'''
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='creates a YOLOv{5,6} data.yaml and trains it')
parser.add_argument('csv', metavar='csv', type=str,
help='csv file', default=defaults.MAIN_CSV_PATH)
parser.add_argument('--version', metavar='version', type=int,
help='yolo version to target', default=6)
parser.add_argument('--data', metavar='data', type=str,
help='data path', default=defaults.DATA_PATH)
args = parser.parse_args()
bcos = entity.read_entities(args.csv)
with open(f'{defaults.DATA_PATH}/data.yaml', 'w') as f:
f.write(gen_data_yaml(bcos, args.data, args.version))
|