Spaces:
Runtime error
Runtime error
File size: 1,234 Bytes
cd1bbe4 1f73e71 cd1bbe4 6d7f6ee e04a33d cd1bbe4 6d7f6ee 94a3eb0 6d7f6ee cd1bbe4 acbdf2a cd1bbe4 1f73e71 bbaa3e0 6d7f6ee 26ef429 bbaa3e0 1f73e71 6d7f6ee |
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 |
import csv
import entity
import argparse
from common import defaults
YOLO_TEMPLATES = {
5: '''
train: %%datapath%%/squares
val: %%datapath%%squares
''',
6: '''
train: %%datapath%%/squares/images
val: %%datapath%%/squares/images
test: %%datapath%%/squares/images
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_TEMPLATES[version].replace('%%datapath%%', datapath)}
nc: {len(bcos.keys())}
names: {names}
'''
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='creates a YOLOv5 data.yaml')
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)
print(gen_data_yaml(bcos, args.data, args.version))
|