dog_train_sys.cpp
8.64 KB
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
//
//#include <fstream>
//#include <iostream>
//#include <sstream>
//#include <vector>
//
//
//#include "NvInfer.h"
//#include "NvOnnxParser.h"
//#include <opencv2/opencv.hpp>
//
//#include "AlgorithmResult.h"
//#include "cuda_kernels.h"
//
//
//// @brief 用于创建IBuilder、IRuntime或IRefitter实例的记录器用于通过该接口创建的所有对象。
//// 在释放所有创建的对象之前，记录器应一直有效。
//// 主要是实例化ILogger类下的log()方法。
//class Logger : public nvinfer1::ILogger
//{
// void log(Severity severity, const char* message) noexcept
// {
// // suppress info-level messages
// if (severity != Severity::kINFO)
// std::cout << message << std::endl;
// }
//} gLogger;
//
//
//
//void onnx_to_engine(std::string onnx_file_path, std::string engine_file_path, int type) {
//
// // 构建器，获取cuda内核目录以获取最快的实现
// // 用于创建config、network、engine的其他对象的核心类
// nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
// const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
// // 解析onnx网络文件
// // tensorRT模型类
// nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
// // onnx文件解析类
// // 将onnx文件解析，并填充rensorRT网络结构
// nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
// // 解析onnx文件
// parser->parseFromFile(onnx_file_path.c_str(), 2);
// for (int i = 0; i < parser->getNbErrors(); ++i) {
// std::cout << "load error: " << parser->getError(i)->desc() << std::endl;
// }
// printf("tensorRT load mask onnx model successfully!!!...\n");
//
// // 创建推理引擎
// // 创建生成器配置对象。
// nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
// // 设置最大工作空间大小。
// config->setMaxWorkspaceSize(16 * (1 << 20));
// // 设置模型输出精度
// if (type == 1) {
// config->setFlag(nvinfer1::BuilderFlag::kFP16);
// }
// if (type == 2) {
// config->setFlag(nvinfer1::BuilderFlag::kINT8);
// }
// // 创建推理引擎
// nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// // 将推理银枪保存到本地
// std::cout << "try to save engine file now~~~" << std::endl;
// std::ofstream file_ptr(engine_file_path, std::ios::binary);
// if (!file_ptr) {
// std::cerr << "could not open plan output file" << std::endl;
// return;
// }
// // 将模型转化为文件流数据
// nvinfer1::IHostMemory* model_stream = engine->serialize();
// // 将文件保存到本地
// file_ptr.write(reinterpret_cast<const char*>(model_stream->data()), model_stream->size());
// // 销毁创建的对象
// model_stream->destroy();
// engine->destroy();
// network->destroy();
// parser->destroy();
// std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl;
//}
//
//int main() {
//
//
// const char* model_path_onnx = "E:/Archime/dog_pose_detect/yolov5/runs/train/exp10/weights/best.onnx";
// const char* model_path_engine = "E:/Archime/dog_pose_detect/yolov5/runs/train/exp10/weights/best.engine";
// const char* image_path = "F:/dog_trainer_sys/test1/IMG_6837.JPG";
// std::string lable_path = "F:/dog_trainer_sys/train2/classes.txt";
// const char* input_node_name = "images";
// const char* output_node_name = "output";
// int num_ionode = 2;
//
// // 读取本地模型文件
// std::ifstream file_ptr(model_path_engine, std::ios::binary);
// if (!file_ptr.good()) {
// std::cerr << "文件无法打开，请确定文件是否可用！" << std::endl;
// }
//
// size_t size = 0;
// file_ptr.seekg(0, file_ptr.end); // 将读指针从文件末尾开始移动0个字节
// size = file_ptr.tellg(); // 返回读指针的位置，此时读指针的位置就是文件的字节数
// file_ptr.seekg(0, file_ptr.beg); // 将读指针从文件开头开始移动0个字节
// char* model_stream = new char[size];
// file_ptr.read(model_stream, size);
// file_ptr.close();
//
// // 日志记录接口
// //Logger logger;
// // 反序列化引擎
// nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
// // 推理引擎
// // 保存模型的模型结构、模型参数以及最优计算kernel配置；
// // 不能跨平台和跨TensorRT版本移植
// nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(model_stream, size);
// // 上下文
// // 储存中间值，实际进行推理的对象
// // 由engine创建，可创建多个对象，进行多推理任务
// nvinfer1::IExecutionContext* context = engine->createExecutionContext();
//
//
// delete[] model_stream;
//
// // 创建GPU显存缓冲区
// void** data_buffer = new void*[num_ionode];
// // 创建GPU显存输入缓冲区
// int input_node_index = engine->getBindingIndex(input_node_name);
// nvinfer1::Dims input_node_dim = engine->getBindingDimensions(input_node_index);
// size_t input_data_length = input_node_dim.d[1] * input_node_dim.d[2] * input_node_dim.d[3];
// cudaMalloc(&(data_buffer[input_node_index]), input_data_length * sizeof(float));
// // 创建GPU显存输出缓冲区
// int output_node_index = engine->getBindingIndex(output_node_name);
// nvinfer1::Dims output_node_dim = engine->getBindingDimensions(output_node_index);
// size_t output_data_length = output_node_dim.d[1] * output_node_dim.d[2];
// cudaMalloc(&(data_buffer[output_node_index]), output_data_length * sizeof(float));
//
//
// // 图象预处理 - 格式化操作
// cv::Mat image = cv::imread(image_path);
// int max_side_length = std::max(image.cols, image.rows);
// cv::Mat max_image = cv::Mat::zeros(cv::Size(max_side_length, max_side_length), CV_8UC3);
// cv::Rect roi(0, 0, image.cols, image.rows);
// image.copyTo(max_image(roi));
// // 将图像归一化，并放缩到指定大小
// cv::Size input_node_shape(input_node_dim.d[2], input_node_dim.d[3]);
// cv::Mat BN_image = cv::dnn::blobFromImage(max_image, 1 / 255.0, input_node_shape, cv::Scalar(0, 0, 0), true, false);
//
// std::vector<float> input_data(input_data_length);
// memcpy(input_data.data(), BN_image.ptr<float>(), input_data_length * sizeof(float));
//
// //void* pGPUData;
// //cudaMalloc(&pGPUData, 3 * image.cols * image.rows * sizeof(unsigned char));
// //cudaMemcpy(pGPUData, (void*)(image.data), 3 * image.cols * image.rows * sizeof(unsigned char), cudaMemcpyHostToDevice);
// //cuda_common::resizeAndNorm((unsigned char*)pGPUData, max_side_length, max_side_length, (float*)data_buffer[input_node_index], input_node_dim.d[2], input_node_dim.d[3]);
//
// // 创建输入cuda流
// cudaStream_t stream;
// cudaStreamCreate(&stream);
//
// // 输入数据由内存到GPU显存
// cudaMemcpyAsync(data_buffer[input_node_index], input_data.data(), input_data_length * sizeof(float), cudaMemcpyHostToDevice, stream);
//
// // 模型推理
// context->enqueueV2(data_buffer, stream, nullptr);
//
// float* result_array = new float[output_data_length];
// cudaMemcpyAsync(result_array, data_buffer[output_node_index], output_data_length * sizeof(float), cudaMemcpyDeviceToHost, stream);
//
// ResultYolov5 result;
// result.factor = max_side_length / (float)input_node_dim.d[2];
// result.read_class_names(lable_path);
//
// //cv::Mat result_image = result.yolov5_result(image, result_array);
//
// //// 查看输出结果
// //cv::imshow("C++ + OpenVINO + Yolov5 推理结果", result_image);
// //cv::waitKey();
//
// std::vector<DogPoseResult> vec_result = result.yolov5_result(result_array, 0.6);
// if (vec_result.size() > 0) {
// DogPoseResult poseResult = vec_result[0];
// std::cout << poseResult.x << std::endl;
// std::cout << poseResult.y << std::endl;
// std::cout << poseResult.width << std::endl;
// std::cout << poseResult.height << std::endl;
// std::cout << poseResult.confidence << std::endl;
// std::cout << poseResult.classId << std::endl;
// std::cout << poseResult.className << std::endl;
//
//
// cv::Rect position_boxe;
// position_boxe.x = poseResult.x;
// position_boxe.y = poseResult.y;
// position_boxe.width = poseResult.width;
// position_boxe.height = poseResult.height;
// cv::rectangle(image, position_boxe, cv::Scalar(0, 0, 255), 2, 8);
// cv::rectangle(image, cv::Point(position_boxe.x, position_boxe.y - 20), cv::Point(position_boxe.x, position_boxe.y), cv::Scalar(0, 255, 255), -1);
// cv::putText(image, poseResult.className, cv::Point(position_boxe.x, position_boxe.y - 10), cv::FONT_HERSHEY_SIMPLEX, .5, cv::Scalar(0, 0, 0));
//
// cv::imwrite("result.jpg", image);
// cv::imshow("show", image);
// cv::waitKey();
// }
//}