write_data.py: add --data arg

Signed-off-by: Niv Sardi <[email protected]>

python/write_data.py

@@ -4,11 +4,13 @@ import argparse
 
 from common import defaults
 
-def gen_data_yaml(bcos):
+def gen_data_yaml(bcos, datapath='../data'):
     names = [f"{d['name']}" for d in bcos.values()]
     return f'''
-train: ../data/squares
-val: ../data/squares
+# this file is autogenerated by write_data.py
+
+train: {datapath}/squares
+val: {datapath}/squares
 
 nc: {len(bcos.keys())}
 names: [{names}]
@@ -18,6 +20,8 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='creates a YOLOv5 data.yaml')
     parser.add_argument('csv', metavar='csv', type=str, 
                     help='csv file')
+    parser.add_argument('--data', metavar='data', type=str,
+                        help='data path', default='../data')
     args = parser.parse_args()
     bcos =  entity.read_entities(args.csv)
-    print(gen_data_yaml(bcos))
+    print(gen_data_yaml(bcos, args.data))