将得到的模型转化为onnx模型,加载到c++中运行,来完成模型的部署,下载并安装onnxruntime;
CMakeLists.txt:
cmake_minimum_required(VERSION 2.8) project(test) #使用clang++编译器 set(CMAKE_CXX_COMPILER clang++) set(CMAKE_BUILD_TYPE "Release") set(CMAKE_INCLUDE_CURRENT_DIR ON) #find the opencv and the qt5 find_package(OpenCV 4.5.1 REQUIRED) #onnxruntime set(ONNXRUNTIME_ROOT_PATH /home/zyl/ubuntu/tensorRT/onnxruntime-master) set(ONNXRUNTIME_INCLUDE_DIRS ${ONNXRUNTIME_ROOT_PATH}/include/onnxruntime ${ONNXRUNTIME_ROOT_PATH}/onnxruntime ${ONNXRUNTIME_ROOT_PATH}/include/onnxruntime/core/session/) set(ONNXRUNTIME_LIB ${ONNXRUNTIME_ROOT_PATH}/build/Linux/Release/libonnxruntime.so) include_directories(${ONNXRUNTIME_INCLUDE_DIRS}) add_executable(${PROJECT_NAME} "main.cpp") target_link_libraries(${PROJECT_NAME} ${OpenCV_LIBS} ${ONNXRUNTIME_LIB})
C++源码:
#include <iostream> #include <opencv2/highgui.hpp> #include <opencv2/imgproc.hpp> #include <opencv2/core.hpp> #include <opencv2/imgcodecs.hpp> #include <chrono> #include <string> //onnxruntime #include <core/session/onnxruntime_cxx_api.h> #include <core/providers/cuda/cuda_provider_factory.h> #include <core/session/onnxruntime_c_api.h> #include <core/providers/tensorrt/tensorrt_provider_factory.h> using namespace std; int main(int argc,char ** argv) { //输入网络的维度 static constexpr const int width = 600; static constexpr const int height = 600; static constexpr const int channel = 3; std::array<int64_t, 4> input_shape_{ 1,height, width,channel}; //-------------------------------------------------------------onnxruntime------------------------------------------------- //图片和模型位置 string img_path = argv[1]; string model_path = "model.onnx"; cv::Mat imgSource = cv::imread(img_path); Ort::Env env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "Detection"); Ort::SessionOptions session_options; //CUDA加速开启 OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0); session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); Ort::AllocatorWithDefaultOptions allocator; //加载ONNX模型 Ort::Session session(env, model_path.c_str(), session_options); //获取输入输出的维度 std::vector<int64_t> input_dims = session.GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); std::vector<int64_t> output_dims = session.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape(); /* session.GetOutputName(1, allocator); session.GetInputName(1, allocator); //输出模型输入节点的数量 size_t num_input_nodes = session.GetInputCount(); size_t num_output_nodes = session.GetOutputCount(); */ std::vector<const char*> input_node_names = {"image_tensor:0"}; std::vector<const char*> output_node_names = {"detection_boxes:0","detection_scores:0","detection_classes:0","num_detections:0" }; input_dims[0] = output_dims[0] = 1;//batch size = 1 std::vector<Ort::Value> input_tensors; auto memory_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); //将图像存储到uchar数组中,BGR--->RGB std::array<uchar, width * height *channel> input_image_{}; uchar* input = input_image_.data(); for (int i = 0; i < imgSource.rows; i++) { for (int j = 0; j < imgSource.cols; j++) { for (int c = 0; c < 3; c++) { //NHWC 格式 if(c==0) input[i*imgSource.cols*3+j*3+c] = imgSource.ptr<uchar>(i)[j*3+2]; if(c==1) input[i*imgSource.cols*3+j*3+c] = imgSource.ptr<uchar>(i)[j*3+1]; if(c==2) input[i*imgSource.cols*3+j*3+c] = imgSource.ptr<uchar>(i)[j*3+0]; //NCHW 格式 // if (c == 0) // input[c*imgSource.rows*imgSource.cols + i * imgSource.cols + j] = imgSource.ptr<uchar>(i)[j * 3 + c]; // if (c == 1) // input[c*imgSource.rows*imgSource.cols + i * imgSource.cols + j] = imgSource.ptr<uchar>(i)[j * 3 + c]; // if (c == 2) // input[c*imgSource.rows*imgSource.cols + i * imgSource.cols + j] = imgSource.ptr<uchar>(i)[j * 3 + c]; } } } input_tensors.push_back(Ort::Value::CreateTensor<uchar>( memory_info, input, input_image_.size(), input_shape_.data(), input_shape_.size())); //不知道输入维度时 //input_tensors.push_back(Ort::Value::CreateTensor<uchar>( // memory_info, input, input_image_.size(), input_dims.data(), input_dims.size())); chrono::steady_clock::time_point t1 = chrono::steady_clock::now(); std::vector<Ort::Value> output_tensors; for(int i=0; i<100;i++) output_tensors = session.Run(Ort::RunOptions { nullptr }, input_node_names.data(), //输入节点名 input_tensors.data(), //input tensors input_tensors.size(), //1 output_node_names.data(), //输出节点名 output_node_names.size()); //4 chrono::steady_clock::time_point t2 = chrono::steady_clock::now(); chrono::duration<double> delay_time = chrono::duration_cast<chrono::duration<double>>(t2 - t1); //milliseconds 毫秒 cout<<"前向传播平均耗时:"<<delay_time.count()*1000/100.0<<"ms"<<endl; float* boxes_ = output_tensors[0].GetTensorMutableData<float>(); float* scores_ = output_tensors[1].GetTensorMutableData<float>(); float* class_ = output_tensors[2].GetTensorMutableData<float>(); float* num_detection = output_tensors[3].GetTensorMutableData<float>(); //-------------------------------------------------------------onnxruntime------------------------------------------------- //------------------循环遍历显示检测框-------------------------- cv::Mat frame(imgSource.clone()); for (int i = 0; i < num_detection[0]; i++) { float confidence = scores_[i]; size_t objectClass = (size_t)class_[i]; if (confidence >= 0.8) { int xLeftBottom = static_cast<int>(boxes_[i*4 + 1] * frame.cols); int yLeftBottom = static_cast<int>(boxes_[i*4 + 0] * frame.rows); int xRightTop = static_cast<int>(boxes_[i*4 + 3] * frame.cols); int yRightTop = static_cast<int>(boxes_[i*4 + 2]* frame.rows); //显示检测框 cv::Rect object((int)xLeftBottom, (int)yLeftBottom, (int)(xRightTop - xLeftBottom), (int)(yRightTop - yLeftBottom)); cv::rectangle(frame, object, cv::Scalar(0,0,255), 2); cv::String label = cv::String("confidence :") +to_string(confidence); int baseLine = 0; cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.3, 1, &baseLine); cv::rectangle(frame, cv::Rect(cv::Point(xLeftBottom, yLeftBottom - labelSize.height), cv::Size(labelSize.width, labelSize.height + baseLine)), cv::Scalar(255, 255, 0), cv::FILLED); cv::putText(frame, label, cv::Point(xLeftBottom, yLeftBottom), cv::FONT_HERSHEY_SIMPLEX, 0.3, cv::Scalar(0, 0, 0)); } } cv::imshow("frame",frame); cv::waitKey(0); return 0; }
参考链接: