Add gen_pkl

gen_pkl.py

@@ -0,0 +1,47 @@
+import argparse
+import pickle
+from collections import OrderedDict
+
+from ultralytics import YOLO
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-w',
+                        '--weights',
+                        type=str,
+                        required=True,
+                        help='YOLOv8 pytorch weights')
+    parser.add_argument('-o',
+                        '--output',
+                        type=str,
+                        required=True,
+                        help='Output file')
+    args = parser.parse_args()
+    return args
+
+
+args = parse_args()
+
+model = YOLO(args.weights)
+model.model.fuse()
+YOLOv8 = model.model.model
+
+strides = YOLOv8[-1].stride.detach().cpu().numpy()
+reg_max = YOLOv8[-1].dfl.conv.weight.shape[1]
+
+state_dict = OrderedDict(GD=model.model.yaml['depth_multiple'],
+                         GW=model.model.yaml['width_multiple'],
+                         strides=strides,
+                         reg_max=reg_max)
+
+for name, value in YOLOv8.state_dict().items():
+    value = value.detach().cpu().numpy()
+    i = int(name.split('.')[0])
+    layer = YOLOv8[i]
+    module_name = layer.type.split('.')[-1]
+    stem = module_name + '.' + name
+    state_dict[stem] = value
+
+with open(args.output, 'wb') as f:
+    pickle.dump(state_dict, f)