-
Notifications
You must be signed in to change notification settings - Fork 0
/
onnx2rknn_step1.py
70 lines (57 loc) · 2.23 KB
/
onnx2rknn_step1.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
import os, glob, shutil
from rknn.api import RKNN
input_width = 640
input_height = 480
model_path = "./model"
dataset_path = "./dataset"
config_path = "./config"
dataset_file = "./dataset.txt"
model_name = 'yolov7-tiny'
platform = "rk3588"
ONNX_MODEL = f'{model_path}/{model_name}-{input_height}-{input_width}.onnx'
OUT_NODE = ["361"]
def get_dataset_txt(dataset_path, dataset_savefile):
file_data = glob.glob(os.path.join(dataset_path,"*.jpg"))
with open(dataset_savefile, "w") as f:
for file in file_data:
f.writelines(f"{file}\n")
def move_onnx_config():
file_data = glob.glob("*.onnx")
for file in file_data:
shutil.move(file, f"{config_path}/{file}")
if __name__ == '__main__':
isExist = os.path.exists(dataset_path)
if not isExist:
os.makedirs(dataset_path)
isExist = os.path.exists(config_path)
if not isExist:
os.makedirs(config_path)
# Prepare the dataset text file
get_dataset_txt(dataset_path, dataset_file)
# Create RKNN object
rknn = RKNN(verbose=False)
# pre-process config
print('--> Config model')
rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform=platform)
print('done')
# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL, outputs=OUT_NODE)
if ret != 0:
print('Load model failed!')
exit(ret)
print('done')
# Build model
print('--> hybrid_quantization_step1')
ret = rknn.hybrid_quantization_step1(dataset=dataset_file, proposal=False)
if ret != 0:
print('hybrid_quantization_step1 failed!')
exit(ret)
print('done')
rknn.release()
print('--> Move hybrid quatization config into config folder')
shutil.move(f"{model_name}-{input_height}-{input_width}.data", f"{config_path}/{model_name}-{input_height}-{input_width}.data")
shutil.move(f"{model_name}-{input_height}-{input_width}.model", f"{config_path}/{model_name}-{input_height}-{input_width}.model")
shutil.move(f"{model_name}-{input_height}-{input_width}.quantization.cfg", f"{config_path}/{model_name}-{input_height}-{input_width}.quantization.cfg")
print('--> Move onnx config into config folder')
move_onnx_config()