python/write_data.py

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