File size: 805 Bytes
cd1bbe4
 
1f73e71
 
cd1bbe4
 
94a3eb0
e04a33d
cd1bbe4
94a3eb0
 
 
 
cd1bbe4
 
acbdf2a
cd1bbe4
 
 
1f73e71
 
bbaa3e0
26ef429
bbaa3e0
1f73e71
 
94a3eb0
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
import csv
import entity
import argparse

from common import defaults

def gen_data_yaml(bcos, datapath='../data'):
    names = [f"{d.name}" for d in bcos.values()]
    return f'''
# this file is autogenerated by write_data.py

train: {datapath}/squares
val: {datapath}/squares

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