Skip to content

Commit 8d1e38c

Browse files
committed
resnet densenet yolo
1 parent 7e42bfc commit 8d1e38c

File tree

9 files changed

+58
-137
lines changed

9 files changed

+58
-137
lines changed

.gitignore

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,6 @@ build
33
build*
44
3rdparty/tensorflow
55
app/AccuracyImgNet/imgs
6-
docs/model_data_alexnet_1.json
6+
docs/
77
docs/input
88
docs/mnist

app/Converters/CMakeLists.txt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,7 @@ add_executable(Reader_weights_onnx reader_weights_sample_onnx.cpp)
33
target_link_libraries(Reader_weights PUBLIC perf_lib layers_lib reader_lib)
44
target_link_libraries(Reader_weights_onnx PUBLIC perf_lib layers_lib reader_lib)
55
add_definitions(-DMODEL_PATH_H5="${CMAKE_SOURCE_DIR}/docs/jsons/model_data_alexnet_1.json")
6-
add_definitions(-DMODEL_PATH_GOOGLENET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/googlenet_onnx_model.json")
6+
add_definitions(-DMODEL_PATH_GOOGLENET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/googlenet_onnx_model.json")
7+
add_definitions(-DMODEL_PATH_DENSENET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/densenet121_Opset16_onnx_model.json")
8+
add_definitions(-DMODEL_PATH_RESNET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/resnest101e_Opset16_onnx_model.json")
9+
add_definitions(-DMODEL_PATH_YOLO11NET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/yolo11x-cls_onnx_model.json")

app/Converters/parser_onnx.py

Lines changed: 41 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,35 @@
11
import json
22
import onnx
33
import os
4+
from onnx import TensorProto
45
from onnx import helper, numpy_helper
6+
from ultralytics import YOLO
7+
8+
def convert_pt_to_onnx(pt_model_path, onnx_model_path=None):
9+
"""Конвертирует YOLO .pt модель в ONNX формат"""
10+
if not pt_model_path.endswith('.pt'):
11+
raise ValueError("Файл модели должен иметь расширение .pt")
12+
13+
if onnx_model_path is None:
14+
onnx_model_path = pt_model_path.replace('.pt', '.onnx')
15+
16+
# Загрузка и экспорт модели
17+
model = YOLO(pt_model_path)
18+
model.export(format="onnx", dynamic=False, simplify=True)
19+
20+
# Проверяем, что файл создан
21+
if not os.path.exists(onnx_model_path):
22+
raise RuntimeError(f"Не удалось создать ONNX файл по пути: {onnx_model_path}")
23+
24+
return onnx_model_path
525

626

727
def onnx_to_json(model_path, output_json_path):
28+
# Проверяем формат модели
29+
if model_path.endswith('.pt'):
30+
print(f"Обнаружена модель .pt, конвертируем в ONNX...")
31+
model_path = convert_pt_to_onnx(model_path)
32+
833
# Загрузка модели
934
model = onnx.load(model_path)
1035
onnx.checker.check_model(model)
@@ -40,7 +65,9 @@ def onnx_to_json(model_path, output_json_path):
4065
# Обработка атрибутов
4166
for attr in node.attribute:
4267
attr_value = helper.get_attribute_value(attr)
43-
if isinstance(attr_value, bytes):
68+
if isinstance(attr_value, TensorProto):
69+
attr_value = numpy_helper.to_array(attr_value).tolist()
70+
elif isinstance(attr_value, bytes):
4471
attr_value = attr_value.decode('utf-8', errors='ignore')
4572
elif hasattr(attr_value, 'tolist'):
4673
attr_value = attr_value.tolist()
@@ -61,17 +88,13 @@ def onnx_to_json(model_path, output_json_path):
6188
if input_name in initializers_dict:
6289
node_init.append(initializers_dict[input_name])
6390

64-
# Новая логика: разделяем weights/value/bias
6591
if len(node_init) == 1:
6692
init = node_init[0]
6793
if len(init["dims"]) == 0 or (len(init["dims"]) == 1 and init["dims"][0] == 1):
68-
# Скалярное значение или массив из одного элемента
6994
layer_data["value"] = init["values"] if len(init["dims"]) == 0 else init["values"][0]
7095
else:
71-
# Многомерные данные
7296
layer_data["weights"] = init["values"]
7397
elif len(node_init) > 1:
74-
# Для нескольких инициализаторов: weights + bias
7598
weights = []
7699
for init in node_init[:-1]:
77100
if len(init["dims"]) > 0:
@@ -81,15 +104,21 @@ def onnx_to_json(model_path, output_json_path):
81104
if weights:
82105
layer_data["weights"] = weights
83106

84-
# Последний инициализатор - bias (если одномерный)
85107
if len(node_init[-1]["dims"]) == 1:
86108
layer_data["bias"] = node_init[-1]["values"]
87109

88110
layer_info.append(layer_data)
89111

90112
class CustomEncoder(json.JSONEncoder):
91113
def default(self, obj):
92-
if hasattr(obj, 'tolist'):
114+
if isinstance(obj, TensorProto):
115+
return {
116+
"name": obj.name,
117+
"dims": list(obj.dims),
118+
"data_type": obj.data_type,
119+
"raw_data": obj.raw_data.hex() if obj.raw_data else None,
120+
}
121+
elif hasattr(obj, 'tolist'):
93122
return obj.tolist()
94123
elif str(type(obj)).endswith("RepeatedScalarContainer'>"):
95124
return list(obj)
@@ -101,7 +130,10 @@ def default(self, obj):
101130
print(f"Модель успешно сохранена в {output_json_path}")
102131

103132

133+
# Пример использования
104134
BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
105-
MODEL_PATH = os.path.join(BASE_DIR, 'docs\\models', 'GoogLeNet.onnx')
106-
MODEL_DATA_PATH = os.path.join(BASE_DIR, 'docs\\jsons', 'googlenet_onnx_model.json')
135+
136+
MODEL_PATH = os.path.join(BASE_DIR, 'docs\\models', 'yolo11x-cls.pt')
137+
MODEL_DATA_PATH = os.path.join(BASE_DIR, 'docs\\jsons', 'yolo11x-cls_onnx_model.json')
138+
107139
onnx_to_json(MODEL_PATH, MODEL_DATA_PATH)

app/Converters/reader_weights_sample_onnx.cpp

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
#include <iostream>
22

3-
#include "Weights_Reader/reader_weights_onnx.hpp"
3+
#include "Weights_Reader/reader_weights.hpp"
44

55
int main() {
6-
std::string json_file = MODEL_PATH_GOOGLENET_ONNX;
7-
json model_data = read_json_onnx(json_file);
6+
std::string json_file = MODEL_PATH_YOLO11NET_ONNX;
7+
json model_data = read_json(json_file);
88

99
std::cout << "Model contains " << model_data.size()
1010
<< " layers:" << std::endl;
@@ -22,7 +22,6 @@ int main() {
2222
std::cout << "Layer " << layer_index << ": " << layer_name << " ("
2323
<< layer_type << ")" << std::endl;
2424

25-
// Вывод атрибутов
2625
if (layer_data.contains("attributes") &&
2726
!layer_data["attributes"].empty()) {
2827
std::cout << " Attributes:" << std::endl;
@@ -46,7 +45,6 @@ int main() {
4645
}
4746
}
4847

49-
// Обработка value (скалярных значений)
5048
if (has_value) {
5149
try {
5250
float value = layer_data["value"].get<float>();
@@ -56,10 +54,9 @@ int main() {
5654
}
5755
}
5856

59-
// Обработка весов
6057
if (has_weights) {
6158
try {
62-
Tensor tensor = create_tensor_from_json_onnx(layer_data, Type::kFloat);
59+
Tensor tensor = create_tensor_from_json(layer_data, Type::kFloat);
6360

6461
std::cout << " Weights shape: " << tensor.get_shape() << std::endl;
6562

app/Graph/CMakeLists.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,4 +72,7 @@ file(DOWNLOAD
7272
add_definitions(-DIMAGE1_PATH="${CMAKE_SOURCE_DIR}/docs/input/")
7373
add_definitions(-DMODEL_PATH_H5="${CMAKE_SOURCE_DIR}/docs/jsons/model_data_alexnet_1.json")
7474
add_definitions(-DMODEL_PATH_GOOGLENET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/googlenet_onnx_model.json")
75+
add_definitions(-DMODEL_PATH_DENSENET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/densenet121_Opset16_onnx_model.json")
76+
add_definitions(-DMODEL_PATH_RESNET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/resnest101e_Opset16_onnx_model.json")
77+
add_definitions(-DMODEL_PATH_YOLO11NET_ONNX="${CMAKE_SOURCE_DIR}/docs/jsons/yolo11x-cls_onnx_model.json")
7578
add_definitions(-DMNIST_PATH="${CMAKE_SOURCE_DIR}/docs/mnist/mnist/test")

include/Weights_Reader/reader_weights_onnx.hpp

Lines changed: 0 additions & 14 deletions
This file was deleted.

src/Weights_Reader/reader_weights.cpp

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,13 +60,11 @@ Tensor create_tensor_from_json(const json& layer_data, Type type) {
6060
extract_values_from_json(layer_data["weights"], weights);
6161
}
6262

63-
// Извлекаем bias (если есть)
6463
std::vector<float> bias;
6564
if (layer_data.contains("bias") && !layer_data["bias"].empty()) {
6665
extract_values_from_json(layer_data["bias"], bias);
6766
}
6867

69-
// Определяем shape
7068
std::vector<size_t> shape;
7169
if (layer_data.contains("weights")) {
7270
parse_json_shape(layer_data["weights"], shape);

src/Weights_Reader/reader_weights_onnx.cpp

Lines changed: 0 additions & 82 deletions
This file was deleted.

test/model_read/model_read.cpp

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -92,39 +92,23 @@ TEST(ParseJsonShapeTests, HandlesNestedArray) {
9292
TEST(ExtractValuesWithoutBiasTest, HandlesCaseWithoutBias) {
9393
json j = json::array({{1.0, 2.0}, {3.0, 4.0}});
9494
std::vector<float> values;
95-
extract_values_without_bias(j, values);
96-
std::vector<float> expected = {1.0, 2.0};
95+
extract_values_from_json(j, values);
96+
std::vector<float> expected = {1.0, 2.0, 3.0, 4.0};
9797
EXPECT_EQ(values, expected);
9898
}
9999

100100
TEST(ExtractValuesWithoutBiasTest, HandlesEmptyJson) {
101101
json j = json::array({});
102102
std::vector<float> values;
103-
extract_values_without_bias(j, values);
103+
extract_values_from_json(j, values);
104104
EXPECT_TRUE(values.empty());
105105
}
106106

107107
TEST(ExtractValuesWithoutBiasTest, HandlesComplexNestedCase) {
108108
json j = json::array({{{1.0, 2.0}, {3.0, 4.0}}, {5.0, 6.0}});
109109
std::vector<float> values;
110-
extract_values_without_bias(j, values);
111-
std::vector<float> expected = {1.0, 2.0, 3.0, 4.0};
112-
EXPECT_EQ(values, expected);
113-
}
114-
115-
TEST(ExtractBiasFromJsonTests, extract_bias) {
116-
json j = json::array({{1.0, 2.0}, {3.0, 4.0}});
117-
std::vector<float> values;
118-
extract_bias_from_json(j, values);
119-
std::vector<float> expected = {3.0, 4.0};
120-
EXPECT_EQ(values, expected);
121-
}
122-
123-
TEST(ExtractBiasFromJsonTests, extract_bias_error_array) {
124-
json j = json::array({{1.0, 2.0}, {3.0, 4.0}});
125-
std::vector<float> values;
126-
extract_bias_from_json(j, values);
127-
std::vector<float> expected = {3.0, 4.0};
110+
extract_values_from_json(j, values);
111+
std::vector<float> expected = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
128112
EXPECT_EQ(values, expected);
129113
}
130114

0 commit comments

Comments
 (0)