-
Notifications
You must be signed in to change notification settings - Fork 4
/
yolov5_ultralytics_to_darknet.py
82 lines (72 loc) · 2.74 KB
/
yolov5_ultralytics_to_darknet.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
# YOLO V5 Ultralytics model and config format to Darknet format
# Adriano Santos.
import torch
import struct
import yaml
import os
from configparser import ConfigParser
# Input files
pt_file_name = "yolov5x.pt"
yaml_file_name = "yolov5x.yaml"
# Output files
cfg_file_name = "yolov5x.cfg"
weights_file_name = "yolov5x.weights"
if os.path.exists(cfg_file_name):
os.remove(cfg_file_name)
if os.path.exists(weights_file_name):
os.remove(weights_file_name)
# Config file conversion (YAML to CFG)
with open(cfg_file_name, 'w') as cfg:
with open(yaml_file_name,'r') as f:
f = yaml.load(f,Loader=yaml.FullLoader)
net = 0
for l in f:
if ((l != 'anchors')and(l !='backbone')and(l !='head')):
if net == 0:
cfg.write('\n[net]\n')
cfg.write("width={}\n".format(416))
cfg.write("height={}\n".format(416))
cfg.write("channels={}\n".format(3))
net =1
cfg.write("{}={}\n".format(l,f[l]))
elif (l == 'anchors'):
anchors = []
for v in f[l]:
anchors.extend(v)
cfg.write("anchors={}\n".format(str(anchors)[1:-1]))
elif((l == 'backbone') or (l =='head')):
for vector in f[l]:
#net type
cfg.write('\n[{}]\n'.format(vector[2]))
#from
if type(vector[0]) == int:
cfg.write("from={}\n".format(str(vector[0])))
elif type(vector[0]) == list:
cfg.write("from={}\n".format(str(vector[0])[1:-1]))
#number
if type(vector[1]) == int:
cfg.write("number={}\n".format(str(vector[1])))
elif type(vector[1]) == list:
cfg.write("number={}\n".format(str(vector[1])[1:-1]))
#params
if type(vector[3]) == int:
cfg.write("args={}\n".format(str(vector[3])))
elif type(vector[3]) == list:
cfg.write("args={}\n".format(str(vector[3])[1:-1]))
with open(cfg_file_name,'r') as f:
for l in f.readlines():
print(l)
# Load model
model = torch.load(pt_file_name, map_location="cpu")['model'].float() # load to FP32
model.eval()
# Model create
with open(weights_file_name,'w') as f:
for k,v in model.state_dict().items():
list_data = list(v.reshape(-1).numpy())
f.write('{}'.format(k))
for value in list_data:
f.write(" ")
f.write(struct.pack('>f',float(value)).hex())
f.write('\n')
# Model informations
print(model)