python/write_data: properly support split and yolo v5/6

python/write_data.py

@@ -4,15 +4,15 @@ import argparse
 
 from common import defaults
 
-YOLO_TEMPLATES = {
+YOLO_DATA_TEMPLATES = {
     5: '''
-train: %%datapath%%/squares
-val: %%datapath%%squares
+train: %%datapath%%/split/train
+val: %%datapath%%/split/val
     ''',
     6: '''
-train: %%datapath%%/squares/images
-val:   %%datapath%%/squares/images
-test:  %%datapath%%/squares/images
+train: %%datapath%%/split/images/train
+val:   %%datapath%%/split/images/val
+test:  %%datapath%%/split/images/test
 
 is_coco: False
     '''
@@ -23,14 +23,14 @@ def gen_data_yaml(bcos, datapath='../data', version=6):
     return f'''
 # this file is autogenerated by write_data.py for YOLO version {version}
 
-{YOLO_TEMPLATES[version].replace('%%datapath%%', datapath)}
+{YOLO_DATA_TEMPLATES[version].replace('%%datapath%%', datapath)}
 
 nc: {len(bcos.keys())}
 names: {names}
 '''
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='creates a YOLOv5 data.yaml')
+    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,
@@ -39,4 +39,5 @@ if __name__ == '__main__':
                         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))
+    with open(f'{defaults.DATA_PATH}/data.yaml', 'w') as f:
+        f.write(gen_data_yaml(bcos, args.data, args.version))