python/write_data: support yolo 5 and 6

python/write_data.py

@@ -4,13 +4,26 @@ import argparse
 
 from common import defaults
 
-def gen_data_yaml(bcos, datapath='../data'):
+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
+# this file is autogenerated by write_data.py for YOLO version {version}
 
-train: {datapath}/squares
-val: {datapath}/squares
+{YOLO_TEMPLATES[version].replace('%%datapath%%', datapath)}
 
 nc: {len(bcos.keys())}
 names: {names}
@@ -20,8 +33,10 @@ 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))
+    print(gen_data_yaml(bcos, args.data, args.version))