python/write_data.py

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))