55from onnx import helper , numpy_helper
66from ultralytics import YOLO
77
8+
89def convert_pt_to_onnx (pt_model_path , onnx_model_path = None ):
910 if onnx_model_path is None :
1011 onnx_model_path = pt_model_path .replace ('.pt' , '.onnx' )
1112
1213 model = YOLO (pt_model_path )
13- model .export (format = "onnx" , dynamic = False , simplify = True )
14+ model .export (format = "onnx" , dynamic = False , simplify = False )
1415
1516 return onnx_model_path
1617
18+
1719def onnx_to_json (model_path , output_json_path ):
1820 if model_path .endswith ('.pt' ):
1921 model_path = convert_pt_to_onnx (model_path )
@@ -31,12 +33,40 @@ def onnx_to_json(model_path, output_json_path):
3133 }
3234
3335 layer_info = []
36+
37+ input_info = {}
38+ for input in model .graph .input :
39+ if input .name in initializers_dict :
40+ continue
41+
42+ shape = []
43+ for dim in input .type .tensor_type .shape .dim :
44+ if dim .HasField ('dim_value' ):
45+ # 0 означает динамическую размерность в ONNX
46+ shape .append (dim .dim_value if dim .dim_value != 0 else - 1 )
47+ elif dim .HasField ('dim_param' ):
48+ # Обрабатываем именованные параметры размерностей
49+ shape .append (- 1 ) # или можно сохранить как строку: dim.dim_param
50+ else :
51+ shape .append (- 1 ) # неизвестная размерность
52+
53+ input_info = {
54+ "name" : input .name ,
55+ "shape" : shape ,
56+ "data_type" : input .type .tensor_type .elem_type
57+ }
58+ break
59+
3460 input_layer = {
3561 "index" : 0 ,
36- "name" : " input_1"
62+ "name" : input_info . get ( "name" , " input_1") ,
3763 "type" : "InputLayer" ,
3864 "weights" : [],
39- "attributes" : {}
65+ "bias" : [],
66+ "attributes" : {
67+ "shape" : input_info .get ("shape" , []),
68+ "data_type" : input_info .get ("data_type" , 1 )
69+ }
4070 }
4171 layer_info .append (input_layer )
4272
@@ -45,9 +75,14 @@ def onnx_to_json(model_path, output_json_path):
4575 "index" : len (layer_info ),
4676 "name" : node .name .replace ('/' , '_' ),
4777 "type" : node .op_type ,
48- "attributes" : {}
78+ "attributes" : {},
79+ "inputs" : []
4980 }
5081
82+ for input_name in node .input :
83+ if input_name not in initializers_dict :
84+ layer_data ["inputs" ].append (input_name .replace ('/' , '_' ))
85+
5186 for attr in node .attribute :
5287 attr_value = helper .get_attribute_value (attr )
5388 if isinstance (attr_value , TensorProto ):
@@ -67,29 +102,44 @@ def onnx_to_json(model_path, output_json_path):
67102 elif attr .name == "strides" :
68103 layer_data ["strides" ] = attr_value
69104
70- node_init = []
71- for input_name in node .input :
72- if input_name in initializers_dict :
73- node_init .append (initializers_dict [input_name ])
74-
75- if len (node_init ) == 1 :
76- init = node_init [0 ]
77- if len (init ["dims" ]) == 0 or (len (init ["dims" ]) == 1 and init ["dims" ][0 ] == 1 ):
78- layer_data ["value" ] = init ["values" ] if len (init ["dims" ]) == 0 else init ["values" ][0 ]
79- else :
80- layer_data ["weights" ] = init ["values" ]
81- elif len (node_init ) > 1 :
82- weights = []
83- for init in node_init [:- 1 ]:
84- if len (init ["dims" ]) > 0 :
85- weights .extend (init ["values" ]) if isinstance (init ["values" ][0 ], list ) else weights .append (
86- init ["values" ])
87-
88- if weights :
89- layer_data ["weights" ] = weights
90-
91- if len (node_init [- 1 ]["dims" ]) == 1 :
92- layer_data ["bias" ] = node_init [- 1 ]["values" ]
105+ if node .op_type == "BatchNormalization" :
106+ bn_params = []
107+ for input_name in node .input :
108+ if input_name in initializers_dict :
109+ bn_params .append (initializers_dict [input_name ])
110+
111+ if len (bn_params ) >= 4 :
112+ layer_data ["scale" ] = bn_params [0 ]["values" ]
113+ layer_data ["bias" ] = bn_params [1 ]["values" ]
114+ layer_data ["mean" ] = bn_params [2 ]["values" ]
115+ layer_data ["var" ] = bn_params [3 ]["values" ]
116+
117+ layer_data ["weights" ] = []
118+
119+ else :
120+ node_init = []
121+ for input_name in node .input :
122+ if input_name in initializers_dict :
123+ node_init .append (initializers_dict [input_name ])
124+
125+ if len (node_init ) == 1 :
126+ init = node_init [0 ]
127+ if len (init ["dims" ]) == 0 or (len (init ["dims" ]) == 1 and init ["dims" ][0 ] == 1 ):
128+ layer_data ["value" ] = init ["values" ] if len (init ["dims" ]) == 0 else init ["values" ][0 ]
129+ else :
130+ layer_data ["weights" ] = init ["values" ]
131+ elif len (node_init ) > 1 :
132+ weights = []
133+ for init in node_init [:- 1 ]:
134+ if len (init ["dims" ]) > 0 :
135+ weights .extend (init ["values" ]) if isinstance (init ["values" ][0 ], list ) else weights .append (
136+ init ["values" ])
137+
138+ if weights :
139+ layer_data ["weights" ] = weights
140+
141+ if len (node_init [- 1 ]["dims" ]) == 1 :
142+ layer_data ["bias" ] = node_init [- 1 ]["values" ]
93143
94144 layer_info .append (layer_data )
95145
@@ -116,7 +166,7 @@ def default(self, obj):
116166
117167BASE_DIR = os .path .dirname (os .path .dirname (os .path .dirname (os .path .abspath (__file__ ))))
118168
119- MODEL_PATH = os .path .join (BASE_DIR , 'docs\\ models' , 'yolo11x-cls.pt ' )
120- MODEL_DATA_PATH = os .path .join (BASE_DIR , 'docs\\ jsons' , 'yolo11x-cls_onnx_model .json' )
169+ MODEL_PATH = os .path .join (BASE_DIR , 'docs\\ models' , 'resnest101e_Opset16.onnx ' )
170+ MODEL_DATA_PATH = os .path .join (BASE_DIR , 'docs\\ jsons' , 'resnest101e_Opset16_onnx_model .json' )
121171
122172onnx_to_json (MODEL_PATH , MODEL_DATA_PATH )
0 commit comments