Improvements (generate .wts file) (wang-xinyu#691)

* Improve output (wts) file name generation

* Named arguments

* Minor: empty out argument

Author: CristiFati

yolov5/gen_wts.py

@@ -1,21 +1,41 @@
-import torch
-import struct
 import sys
+import argparse
+import os
+import struct
+import torch
 from utils.torch_utils import select_device
 
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
+    parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
+    parser.add_argument('-o', '--output', help='Output (.wts) file path (optional)')
+    args = parser.parse_args()
+    if not os.path.isfile(args.weights):
+        raise SystemExit('Invalid input file')
+    if not args.output:
+        args.output = os.path.splitext(args.weights)[0] + '.wts'
+    elif os.path.isdir(args.output):
+        args.output = os.path.join(
+            args.output,
+            os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
+    return args.weights, args.output
+
+
+pt_file, wts_file = parse_args()
+
 # Initialize
 device = select_device('cpu')
-pt_file = sys.argv[1]
 # Load model
 model = torch.load(pt_file, map_location=device)['model'].float()  # load to FP32
 model.to(device).eval()
 
-with open(pt_file.split('.')[0] + '.wts', 'w') as f:
+with open(wts_file, 'w') as f:
     f.write('{}\n'.format(len(model.state_dict().keys())))
     for k, v in model.state_dict().items():
         vr = v.reshape(-1).cpu().numpy()
         f.write('{} {} '.format(k, len(vr)))
         for vv in vr:
             f.write(' ')
-            f.write(struct.pack('>f',float(vv)).hex())
+            f.write(struct.pack('>f' ,float(vv)).hex())
         f.write('\n')