├── tensorrt_cmake ├── test.sh ├── README.md ├── CMakeLists.txt └── src │ ├── yolov5_trt.h │ └── yolov5_trt.cpp ├── README.md ├── tensorrt ├── yolov5_trt.h └── yolov5_trt.cpp ├── opencv_det └── yolov5_cv.cpp └── onnxruntime └── ort.cpp /tensorrt_cmake/test.sh: -------------------------------------------------------------------------------- 1 | # 最好用绝对路径 2 | rm -r /xxx/build 3 | mkdir /xxx/build 4 | rm -r /xxx/bin 5 | mkdir /xxx/bin 6 | cd /xxx/build 7 | cmake .. 8 | make 9 | /xxx/yolov5_trt 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # yolov5进行c++推理部署 2 | 3 | 使用opencv，onnxruntime和tensorrt部署编译的时候，记得导入相应的头文件和依赖库 4 | 5 | 这里添加使用cmake的方法，也能更好的将一些所需导入的头文件目录和链接库文件展示出来 6 | 以tensorrt为例 7 | 8 | 9 | 10 | 11 | 12 | -------------------------------------------------------------------------------- /tensorrt_cmake/README.md: -------------------------------------------------------------------------------- 1 | # 使用cmake编译步骤 2 | + 将一些头文件，要链接的依赖库用指令写在CMakeLists.txt里，这里都给了注释 3 | + 新建build和bin文件夹 4 | ```cpp 5 | mkdir build 6 | mkdir bin 7 | cd /xxx/build 8 | cmake .. 9 | make 10 | cd /xxx/bin 11 | ./yolov5_trt 12 | ``` 13 | + 最好写一个脚本一套流程下来直接执行生成的执行文件，因为cmake生成之后如果对代码进行了更改要将以前生成的CmakeCache.txt缓存删掉重新编译才行。 14 | -------------------------------------------------------------------------------- /tensorrt_cmake/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | # # ONNX 3 | # # requirement of cmake version 4 | # cmake_minimum_required(VERSION 3.5) 5 | # # project name 6 | # PROJECT(onnx CXX) 7 | 8 | # # set the directory of executable files 9 | # set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin) # 输出文件夹路径 10 | # set(CMAKE_BUILD_TYPE "Release") # 编译类型 11 | 12 | 13 | # # find required opencv 14 | # find_package(OpenCV REQUIRED) # 使用find_package获取opencv的相关信息,如lib、include、version等 15 | # # directory of opencv headers 16 | # include_directories(${OpenCV_INCLUDE_DIRS}) # 头文件夹 17 | 18 | # #onnxruntime 19 | # set(ONNXRUNTIME_ROOT_PATH /data1/yzycode/opencv/onnxruntime-linux-x64-1.6.0/) # 根目录 20 | # set(ONNXRUNTIME_INCLUDE_DIRS ${ONNXRUNTIME_ROOT_PATH}/include/) # 头文件夹 21 | # set(ONNXRUNTIME_LIB ${ONNXRUNTIME_ROOT_PATH}lib/libonnxruntime.so) # 依赖库 22 | 23 | # # name of executable file and path of source file 24 | # add_executable(onnx src/ort.cpp) 25 | # # directory of opencv library 26 | # link_directories(${OpenCV_LIBRARY_DIRS}) # 指定第三方库所在路径,该路径添加到第三方库搜索路径中 27 | # # opencv libraries 28 | # target_link_libraries(onnx ${OpenCV_LIBS}) # 相对路径 29 | 30 | # include_directories(${ONNXRUNTIME_INCLUDE_DIRS}) # 头文件 31 | # target_link_libraries(onnx ${ONNXRUNTIME_LIB}) # 依赖 32 | 33 | 34 | # TRT 35 | # requirement of cmake version 36 | cmake_minimum_required(VERSION 3.5) 37 | # project name 38 | PROJECT(yolov5_trt CXX) 39 | set(CMAKE_CXX_STANDARD 17) 40 | # set the directory of executable files 41 | set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin) # 输出文件夹路径 42 | set(CMAKE_BUILD_TYPE "Release") # 编译类型 43 | 44 | 45 | # find required opencv 46 | find_package(OpenCV REQUIRED) # 使用find_package获取opencv的相关信息,如lib、include、version等 47 | find_package(glog REQUIRED) 48 | 49 | # tensorrt 50 | set(TENSORRT_ROOT_PATH /opt/TensorRT-8.0.1.6/) # 根目录 51 | set(TENSORRT_INCLUDE_DIRS ${TENSORRT_ROOT_PATH}/include/) # 头文件夹 52 | set(TENSORRT_LIB ${TENSORRT_ROOT_PATH}lib/libnvinfer.so ${TENSORRT_ROOT_PATH}lib/libnvonnxparser.so) # 依赖动态库 53 | 54 | 55 | # cuda 56 | set(CUDA_ROOT_PATH /usr/local/cuda-10.2/) # 根目录 57 | set(CUDA_INCLUDE_DIRS ${CUDA_ROOT_PATH}/include/) # 头文件夹 58 | set(CUDA_LIB ${CUDA_ROOT_PATH}lib64/libcudart.so) # 依赖动态库 59 | 60 | # name of executable file and path of source file 61 | add_executable(yolov5_trt src/yolov5_trt.cpp) 62 | 63 | # directory of opencv headers 64 | include_directories(${OpenCV_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIRS} ${CUDA_INCLUDE_DIRS}) # 头文件夹 65 | # directory of opencv library 66 | link_directories(${OpenCV_LIBRARY_DIRS}) # 指定第三方库所在路径,该路径添加到第三方库搜索路径中 67 | # opencv libraries 68 | target_link_libraries(yolov5_trt ${OpenCV_LIBS} ${TENSORRT_LIB} ${CUDA_LIB} ${glog_LIB}) # 相对路径 69 | -------------------------------------------------------------------------------- /tensorrt/yolov5_trt.h: -------------------------------------------------------------------------------- 1 | // yolov5使用tensorrt进行部署的头文件 2 | 3 | #ifndef YOLOV5_TRT_H 4 | #define YOLOV5_TRT_H 5 | 6 | #include 7 | #include 8 | #include // nvidia加载模型进行推理的插件 9 | #include 10 | #include 11 | 12 | 13 | // 自定义配置结构 14 | struct Configuration 15 | { 16 | float confThreshold; // Confidence threshold 17 | float nmsThreshold; // Non-maximum suppression threshold 18 | float objThreshold; //Object Confidence threshold 19 | std::string modelpath; 20 | }; 21 | 22 | // 定义BoxInfo结构类型 23 | typedef struct BoxInfo 24 | { 25 | float x1; 26 | float y1; 27 | float x2; 28 | float y2; 29 | float score; 30 | int label; 31 | } BoxInfo; 32 | 33 | 34 | class YOLOv5 35 | { 36 | public: 37 | YOLOv5(Configuration config); 38 | ~YOLOv5(); 39 | void UnInit(); 40 | void detect(cv::Mat& frame); 41 | private: 42 | float confThreshold; 43 | float nmsThreshold; 44 | float objThreshold; 45 | int inpWidth; 46 | int inpHeight; 47 | std::string classes[80] = {"person", "bicycle", "car", "motorbike", "aeroplane", "bus", 48 | "train", "truck", "boat", "traffic light", "fire hydrant", 49 | "stop sign", "parking meter", "bench", "bird", "cat", "dog", 50 | "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", 51 | "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 52 | "skis", "snowboard", "sports ball", "kite", "baseball bat", 53 | "baseball glove", "skateboard", "surfboard", "tennis racket", 54 | "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", 55 | "banana", "apple", "sandwich", "orange", "broccoli", "carrot", 56 | "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", 57 | "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", 58 | "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", 59 | "sink", "refrigerator", "book", "clock", "vase", "scissors", 60 | "teddy bear", "hair drier", "toothbrush"}; 61 | 62 | const bool keep_ratio = true; 63 | void normalize_(cv::Mat img); // 归一化函数 64 | void nms(std::vector& input_boxes); 65 | cv::Mat resize_image(cv::Mat srcimg, int *newh, int *neww, int *top, int *left); 66 | 67 | void loadOnnx(const std::string strName); 68 | void loadTrt(const std::string strName); 69 | 70 | nvinfer1::ICudaEngine *m_CudaEngine; 71 | nvinfer1::IRuntime *m_CudaRuntime; 72 | nvinfer1::IExecutionContext *m_CudaContext; 73 | cudaStream_t m_CudaStream; // //初始化流,CUDA流的类型为cudaStream_t 74 | int m_iInputIndex; 75 | int m_iOutputIndex; 76 | int m_iClassNums; 77 | int m_iBoxNums; 78 | cv::Size m_InputSize; 79 | void* m_ArrayDevMemory[2]{0}; 80 | void* m_ArrayHostMemory[2]{0}; 81 | int m_ArraySize[2]{0}; 82 | std::vector m_InputWrappers{}; 83 | }; 84 | 85 | 86 | #endif -------------------------------------------------------------------------------- /tensorrt_cmake/src/yolov5_trt.h: -------------------------------------------------------------------------------- 1 | // yolov5使用tensorrt进行部署的头文件 2 | 3 | #ifndef YOLOV5_TRT_H 4 | #define YOLOV5_TRT_H 5 | 6 | #include 7 | #include 8 | #include // nvidia加载模型进行推理的插件 9 | #include 10 | #include 11 | 12 | 13 | // 自定义配置结构 14 | struct Configuration 15 | { 16 | float confThreshold; // Confidence threshold 17 | float nmsThreshold; // Non-maximum suppression threshold 18 | float objThreshold; //Object Confidence threshold 19 | std::string modelpath; 20 | }; 21 | 22 | // 定义BoxInfo结构类型 23 | typedef struct BoxInfo 24 | { 25 | float x1; 26 | float y1; 27 | float x2; 28 | float y2; 29 | float score; 30 | int label; 31 | } BoxInfo; 32 | 33 | 34 | class YOLOv5 35 | { 36 | public: 37 | YOLOv5(Configuration config); 38 | ~YOLOv5(); 39 | void UnInit(); 40 | void detect(cv::Mat& frame); 41 | private: 42 | float confThreshold; 43 | float nmsThreshold; 44 | float objThreshold; 45 | int inpWidth; 46 | int inpHeight; 47 | std::string classes[80] = {"person", "bicycle", "car", "motorbike", "aeroplane", "bus", 48 | "train", "truck", "boat", "traffic light", "fire hydrant", 49 | "stop sign", "parking meter", "bench", "bird", "cat", "dog", 50 | "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", 51 | "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 52 | "skis", "snowboard", "sports ball", "kite", "baseball bat", 53 | "baseball glove", "skateboard", "surfboard", "tennis racket", 54 | "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", 55 | "banana", "apple", "sandwich", "orange", "broccoli", "carrot", 56 | "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", 57 | "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", 58 | "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", 59 | "sink", "refrigerator", "book", "clock", "vase", "scissors", 60 | "teddy bear", "hair drier", "toothbrush"}; 61 | 62 | const bool keep_ratio = true; 63 | void normalize_(cv::Mat img); // 归一化函数 64 | void nms(std::vector& input_boxes); 65 | cv::Mat resize_image(cv::Mat srcimg, int *newh, int *neww, int *top, int *left); 66 | 67 | void loadOnnx(const std::string strName); 68 | void loadTrt(const std::string strName); 69 | 70 | nvinfer1::ICudaEngine *m_CudaEngine; 71 | nvinfer1::IRuntime *m_CudaRuntime; 72 | nvinfer1::IExecutionContext *m_CudaContext; 73 | cudaStream_t m_CudaStream; // //初始化流,CUDA流的类型为cudaStream_t 74 | int m_iInputIndex; 75 | int m_iOutputIndex; 76 | int m_iClassNums; 77 | int m_iBoxNums; 78 | cv::Size m_InputSize; 79 | void* m_ArrayDevMemory[2]{0}; 80 | void* m_ArrayHostMemory[2]{0}; 81 | int m_ArraySize[2]{0}; 82 | std::vector m_InputWrappers{}; 83 | }; 84 | 85 | 86 | #endif -------------------------------------------------------------------------------- /opencv_det/yolov5_cv.cpp: -------------------------------------------------------------------------------- 1 | // yzy 2 | #include 3 | #include 4 | #include 5 | #include // 深度学习模块 6 | #include // 图像处理模块 7 | #include // 高层GUI图形用户界面 8 | 9 | #include 10 | 11 | 12 | // 命名空间 13 | using namespace cv; 14 | using namespace dnn; 15 | using namespace std; 16 | 17 | // 自定义配置结构 18 | struct Configuration 19 | { 20 | public: 21 | float confThreshold; // Confidence threshold 22 | float nmsThreshold; // Non-maximum suppression threshold 23 | float objThreshold; //Object Confidence threshold 24 | string modelpath; 25 | }; 26 | 27 | // 模型 28 | class YOLOv5 29 | { 30 | public: 31 | // 初始化 32 | YOLOv5(Configuration config,bool isCuda); 33 | void detect(Mat& frame); // 检测函数 34 | private: 35 | float confThreshold; 36 | float nmsThreshold; 37 | float objThreshold; 38 | int inpWidth; 39 | int inpHeight; 40 | int num_classes; 41 | string classes[80] = {"person", "bicycle", "car", "motorbike", "aeroplane", "bus", 42 | "train", "truck", "boat", "traffic light", "fire hydrant", 43 | "stop sign", "parking meter", "bench", "bird", "cat", "dog", 44 | "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", 45 | "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 46 | "skis", "snowboard", "sports ball", "kite", "baseball bat", 47 | "baseball glove", "skateboard", "surfboard", "tennis racket", 48 | "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", 49 | "banana", "apple", "sandwich", "orange", "broccoli", "carrot", 50 | "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", 51 | "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", 52 | "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", 53 | "sink", "refrigerator", "book", "clock", "vase", "scissors", 54 | "teddy bear", "hair drier", "toothbrush"}; 55 | //vector class_names; 56 | const bool keep_ratio = true; 57 | Net net; // dnn里的 58 | void drawPred(float conf, int left, int top, int right, int bottom, Mat& frame, int classid); 59 | Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left); 60 | }; 61 | 62 | // 初始化 63 | YOLOv5::YOLOv5(Configuration config,bool isCuda=false) 64 | { 65 | this->confThreshold = config.confThreshold; 66 | this->nmsThreshold = config.nmsThreshold; 67 | this->objThreshold = config.objThreshold; 68 | 69 | this->net = readNet(config.modelpath); // 解析模型onnx权重。dnn.hpp 70 | //cuda // https://blog.csdn.net/cxyhjl/article/details/125383555 71 | if (isCuda) { 72 | net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA); 73 | net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); 74 | cout<<"cuda"<num_classes = sizeof(this->classes)/sizeof(this->classes[0]); // 类别数量 82 | this->inpHeight = 640; 83 | this->inpWidth = 640; 84 | } 85 | 86 | Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left) 87 | { 88 | int srch = srcimg.rows, srcw = srcimg.cols; // 输入高宽 89 | *newh = this->inpHeight; // 指针变量指向输入yolo模型的宽高 90 | *neww = this->inpWidth; 91 | Mat dstimg; // 定义一个目标源 92 | if (this->keep_ratio && srch != srcw) { // 高宽不等 93 | float hw_scale = (float)srch / srcw; // 保存比列 94 | if (hw_scale > 1) { // 按照yolov5的预处理进行处理 95 | *newh = this->inpHeight; 96 | *neww = int(this->inpWidth / hw_scale); // 97 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 98 | *left = int((this->inpWidth - *neww) * 0.5); 99 | // 和yolov5的处理对应,没有进行32的取模运算,这个是用114像素填充到(640,640)了,最后输入还是640,640 100 | copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114); 101 | } 102 | else { 103 | *newh = (int)this->inpHeight * hw_scale; 104 | *neww = this->inpWidth; 105 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 106 | *top = (int)(this->inpHeight - *newh) * 0.5; 107 | copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114); 108 | } 109 | } 110 | else { 111 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 112 | } 113 | return dstimg; 114 | } 115 | 116 | void YOLOv5::drawPred(float conf, int left, int top, int right, int bottom, Mat& frame, int classid) // Draw the predicted bounding box 117 | { 118 | //Draw a rectangle displaying the bounding box 119 | rectangle(frame, Point(left, top), Point(right, bottom), Scalar(0, 0, 255), 2); 120 | 121 | //Get the label for the class name and its confidence 122 | string label = format("%.2f", conf); 123 | label = this->classes[classid] + ":" + label; 124 | 125 | //Display the label at the top of the bounding box 126 | int baseLine; 127 | Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine); 128 | top = max(top, labelSize.height); 129 | //rectangle(frame, Point(left, top - int(1.5 * labelSize.height)), Point(left + int(1.5 * labelSize.width), top + baseLine), Scalar(0, 255, 0), FILLED); 130 | putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1); 131 | } 132 | 133 | // 预测 134 | void YOLOv5::detect(Mat& frame) 135 | { 136 | int newh = 0, neww = 0, padh = 0, padw = 0; 137 | Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw); // 预处理 138 | Mat blob = blobFromImage(dstimg, 1 / 255.0, Size(this->inpWidth, this->inpHeight), Scalar(0, 0, 0), true, false); // return:4-dimensional Mat with NCHW dimensions order. 139 | this->net.setInput(blob); // 设置输出 140 | vector outs; // 要给空的走一遍 141 | this->net.forward(outs, this->net.getUnconnectedOutLayersNames()); // [b,num_pre,(5+classes)] 142 | 143 | int num_proposal = outs[0].size[1]; // 25200 144 | int out_dim2 = outs[0].size[2]; // 145 | if (outs[0].dims > 2) 146 | { 147 | outs[0] = outs[0].reshape(0, num_proposal); // 一般都会大于二维的，所以展平二维[b,num_pre*(5+classes)] 148 | } 149 | /////generate proposals 150 | vector confidences; 151 | vector boxes; // opencv里保存box的 152 | vector classIds; 153 | float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww; 154 | float* pdata = (float*)outs[0].data; // 定义浮点型指针， 155 | for(int i = 0; i < num_proposal; ++i) // 遍历所有的num_pre_boxes 156 | { 157 | int index = i * out_dim2; // prob[b*num_pred_boxes*(classes+5)] 158 | float obj_conf = pdata[index + 4]; // 置信度分数 159 | if (obj_conf > this->objThreshold) // 大于阈值 160 | { 161 | //Mat scores = outs[0].row(row_ind).colRange(5, nout); // 相当于python里的切片操作，每类的预测类别分数 162 | Mat scores(1, this->num_classes, CV_32FC1, pdata+index + 5); // 这样操作更好理解 163 | Point classIdPoint; //定义点 164 | double max_class_socre; // 定义一个double类型的变量保存预测中类别分数最大值 165 | // Get the value and location of the maximum score 166 | minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint); // 求每类类别分数最大的值和索引 167 | max_class_socre *= obj_conf; // 最大的类别分数*置信度 168 | if (max_class_socre > this->confThreshold) // 再次筛选 169 | { 170 | const int class_idx = classIdPoint.x; // 类别索引,在yolo里就是表示第几类 171 | // float cx = (pdata[0] * 2.f - 0.5f + j) * stride; ///cx,映射回原图。对应yolov5里后处理部分 172 | // float cy = (pdata[1] * 2.f - 0.5f + i) * stride; ///cy，但是现在的yolov5导出的onnx已经将这个处理放在里面了 173 | // float w = powf(pdata[2] * 2.f, 2.f) * anchor_w; ///w，所以这里不需要后处理一下了，直接取输出就行 174 | // float h = powf(pdata[3] * 2.f, 2.f) * anchor_h; ///h 175 | 176 | // 经过后处理的只需要直接取就行 177 | float cx = pdata[index]; //x 178 | float cy = pdata[index+1]; //y 179 | float w = pdata[index+2]; //w 180 | float h = pdata[index+3]; //h 181 | 182 | int left = int((cx - padw - 0.5 * w)*ratiow); // *ratiow，变回原图尺寸 183 | int top = int((cy - padh - 0.5 * h)*ratioh); 184 | 185 | confidences.push_back((float)max_class_socre); 186 | boxes.push_back(Rect(left, top, (int)(w*ratiow), (int)(h*ratioh))); //（x,y,w,h） 187 | classIds.push_back(class_idx); // 188 | } 189 | } 190 | 191 | } 192 | 193 | vector indices; 194 | dnn::NMSBoxes(boxes, confidences, this->confThreshold, this->nmsThreshold, indices); 195 | for (size_t i = 0; i < indices.size(); ++i) 196 | { 197 | int idx = indices[i]; 198 | Rect box = boxes[idx]; 199 | this->drawPred(confidences[idx], box.x, box.y, 200 | box.x + box.width, box.y + box.height, frame, classIds[idx]); 201 | } 202 | } 203 | 204 | int main() 205 | { 206 | double timeStart = (double)getTickCount(); 207 | clock_t startTime,endTime; //计算时间 208 | // 自己定义的yolo一些配置超参 209 | Configuration yolo_nets = { 0.3, 0.5, 0.3, "yolov5s.onnx" }; 210 | YOLOv5 yolo_model(yolo_nets,false); 211 | string imgpath = "bus.jpg"; 212 | Mat srcimg = imread(imgpath); 213 | 214 | double nTime = ((double)getTickCount() - timeStart) / getTickFrequency(); 215 | startTime = clock();//计时开始 216 | yolo_model.detect(srcimg); 217 | endTime = clock();//计时结束 218 | cout << "clock_running time is:" <<(double)(endTime - startTime) / CLOCKS_PER_SEC << "s" << endl; 219 | cout << "The whole run time is:" << (double)clock() /CLOCKS_PER_SEC<< "s" << endl; 220 | cout << "getTickCount_running time :" << nTime << "s" << endl; 221 | 222 | static const string kWinName = "Deep learning object detection in OpenCV"; 223 | namedWindow(kWinName, WINDOW_NORMAL); // 自适应调节窗口大小 224 | imwrite("restult_cpu.jpg",srcimg); 225 | imshow(kWinName, srcimg); 226 | waitKey(0); 227 | destroyAllWindows(); 228 | return 0; 229 | } 230 | 231 | 232 | 233 | -------------------------------------------------------------------------------- /tensorrt/yolov5_trt.cpp: -------------------------------------------------------------------------------- 1 | // yzy 2 | // yolov5进行tensorrt部署的源文件 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "yolov5_trt.h" 8 | 9 | 10 | // 命名空间 11 | using namespace cv; 12 | using namespace nvinfer1; 13 | 14 | 15 | // Logger for TRT info/warning/errors, https://github.com/onnx/onnx-tensorrt/blob/main/onnx_trt_backend.cpp 16 | class TRT_Logger : public nvinfer1::ILogger 17 | { 18 | nvinfer1::ILogger::Severity _verbosity; 19 | std::ostream* _ostream; 20 | 21 | public: 22 | TRT_Logger(Severity verbosity = Severity::kWARNING, std::ostream& ostream = std::cout) 23 | : _verbosity(verbosity) 24 | , _ostream(&ostream) 25 | { 26 | } 27 | void log(Severity severity, const char* msg) noexcept override 28 | { 29 | if (severity <= _verbosity) 30 | { 31 | time_t rawtime = std::time(0); 32 | char buf[256]; 33 | strftime(&buf[0], 256, "%Y-%m-%d %H:%M:%S", std::gmtime(&rawtime)); 34 | const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? " BUG" : severity == Severity::kERROR 35 | ? " ERROR" 36 | : severity == Severity::kWARNING ? "WARNING" : severity == Severity::kINFO ? " INFO" 37 | : "UNKNOWN"); 38 | (*_ostream) << "[" << buf << " " << sevstr << "] " << msg << std::endl; 39 | } 40 | } 41 | }; 42 | 43 | // 判断文件是否形成 44 | static bool ifFileExists(const char *FileName) 45 | { 46 | struct stat my_stat; 47 | return (stat(FileName, &my_stat) == 0); 48 | } 49 | 50 | 51 | // 加载onnx文件 52 | void YOLOv5::loadOnnx(const std::string strModelName) 53 | { 54 | TRT_Logger gLogger; // 日志 55 | //根据tensorrt pipeline 构建网络 56 | IBuilder* builder = createInferBuilder(gLogger); // 网络元数据,用于搭建网络入口 57 | builder->setMaxBatchSize(1); // batchsize 58 | const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); // 显式批处理 59 | INetworkDefinition* network = builder->createNetworkV2(explicitBatch); // 定义模型 60 | nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger); // 使用nvonnxparser 定义一个可用的onnx解析器 61 | parser->parseFromFile(strModelName.c_str(), static_cast(ILogger::Severity::kWARNING)); // 解析onnx 62 | // 使用builder对象构建engine 63 | IBuilderConfig* config = builder->createBuilderConfig(); // 64 | // 特别重要的属性是最大工作空间大小 65 | config->setMaxWorkspaceSize(1ULL << 30); // 分配内存空间 66 | m_CudaEngine = builder->buildEngineWithConfig(*network, *config); // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重 67 | 68 | std::string strTrtName = strModelName; 69 | size_t sep_pos = strTrtName.find_last_of("."); 70 | strTrtName = strTrtName.substr(0, sep_pos) + ".trt"; // 71 | IHostMemory *gieModelStream = m_CudaEngine->serialize(); // 将引擎序列化 72 | std::string serialize_str; // 73 | std::ofstream serialize_output_stream; 74 | serialize_str.resize(gieModelStream->size()); 75 | // memcpy内存拷贝函数 ，从源内存地址的起始位置开始拷贝若干个字节到目标内存地址中 76 | memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size()); 77 | serialize_output_stream.open(strTrtName.c_str()); 78 | serialize_output_stream<createExecutionContext(); //执行上下文用于执行推理 81 | // 使用一次，销毁parser，network, builder, and config 82 | parser->destroy(); 83 | network->destroy(); 84 | config->destroy(); 85 | builder->destroy(); 86 | } 87 | 88 | 89 | void YOLOv5::loadTrt(const std::string strName) 90 | { 91 | TRT_Logger gLogger; 92 | // 序列化引擎被保留并保存到文件中 93 | m_CudaRuntime = createInferRuntime(gLogger); 94 | std::ifstream fin(strName); 95 | std::string cached_engine = ""; 96 | while (fin.peek() != EOF) 97 | { 98 | std::stringstream buffer; 99 | buffer << fin.rdbuf(); 100 | cached_engine.append(buffer.str()); 101 | } 102 | fin.close(); 103 | m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr); // runtime对象反序列化 104 | m_CudaContext = m_CudaEngine->createExecutionContext(); //可以查询引擎获取有关网络的输入和输出的张量信息--维度/数据格式/数据类型 105 | m_CudaRuntime->destroy(); 106 | } 107 | 108 | 109 | // 初始化 110 | YOLOv5::YOLOv5(Configuration config) 111 | { 112 | confThreshold = config.confThreshold; 113 | nmsThreshold = config.nmsThreshold; 114 | objThreshold = config.objThreshold; 115 | inpHeight = 640; 116 | inpWidth = 640; 117 | 118 | std::string model_path = config.modelpath; // 模型权重路径 119 | // 加载模型 120 | std::string strTrtName = config.modelpath; // 加载模型权重 121 | size_t sep_pos = model_path.find_last_of("."); 122 | strTrtName = model_path.substr(0, sep_pos) + ".engine"; // ".trt" 123 | if(ifFileExists(strTrtName.c_str())) 124 | { 125 | loadTrt(strTrtName); 126 | } 127 | else 128 | { 129 | loadOnnx(config.modelpath); 130 | } 131 | 132 | // 利用加载的模型获取输入输出信息 133 | // 使用输入和输出blob名来获取输入和输出索引 134 | m_iInputIndex = m_CudaEngine->getBindingIndex("images"); // 输入索引 135 | m_iOutputIndex = m_CudaEngine->getBindingIndex("output"); // 输出 136 | Dims dims_i = m_CudaEngine->getBindingDimensions(m_iInputIndex); // 输入， 137 | int size1 = dims_i.d[0] * dims_i.d[1] * dims_i.d[2] * dims_i.d[3]; // 展平 138 | m_InputSize = cv::Size(dims_i.d[3], dims_i.d[2]); // 输入尺寸(W,H) 139 | Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex); // 输出，维度[0,1,2,3]NHWC 140 | int size2 = dims_o.d[0] * dims_o.d[1] * dims_o.d[2]; // 所有大小 141 | m_iClassNums = dims_o.d[2] - 5; // [,,classes+5] 142 | m_iBoxNums = dims_o.d[1]; // [b,num_pre_boxes,classes+5] 143 | 144 | // 分配内存大小 145 | cudaMalloc(&m_ArrayDevMemory[m_iInputIndex], size1 * sizeof(float)); 146 | m_ArrayHostMemory[m_iInputIndex] = malloc(size1 * sizeof(float)); 147 | m_ArraySize[m_iInputIndex] = size1 *sizeof(float); 148 | cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float)); 149 | m_ArrayHostMemory[m_iOutputIndex] = malloc( size2 * sizeof(float)); 150 | m_ArraySize[m_iOutputIndex] = size2 *sizeof(float); 151 | 152 | // bgr 153 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex]); 154 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + sizeof(float) * dims_i.d[2] * dims_i.d[3] ); 155 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + 2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]); 156 | //创建CUDA流,推理时TensorRT执行通常是异步的，因此将内核排入CUDA流 157 | cudaStreamCreate(&m_CudaStream); // 只需初始化一次即可 158 | } 159 | 160 | void YOLOv5::UnInit() 161 | { 162 | 163 | for(auto &p: m_ArrayDevMemory) 164 | { 165 | cudaFree(p); 166 | p = nullptr; 167 | } 168 | for(auto &p: m_ArrayHostMemory) 169 | { 170 | free(p); 171 | p = nullptr; 172 | } 173 | cudaStreamDestroy(m_CudaStream); 174 | //m_CudaContext->destroy(); // 这个报错 175 | m_CudaEngine->destroy(); 176 | 177 | } 178 | 179 | YOLOv5::~YOLOv5() 180 | { 181 | UnInit(); 182 | } 183 | 184 | 185 | 186 | Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left) 187 | { 188 | int srch = srcimg.rows, srcw = srcimg.cols; 189 | *newh = this->inpHeight; 190 | *neww = this->inpWidth; 191 | Mat dstimg; 192 | if (this->keep_ratio && srch != srcw) { 193 | float hw_scale = (float)srch / srcw; 194 | if (hw_scale > 1) { 195 | *newh = this->inpHeight; 196 | *neww = int(this->inpWidth / hw_scale); 197 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 198 | *left = int((this->inpWidth - *neww) * 0.5); 199 | copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114); 200 | } 201 | else { 202 | *newh = (int)this->inpHeight * hw_scale; 203 | *neww = this->inpWidth; 204 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 205 | *top = (int)(this->inpHeight - *newh) * 0.5; 206 | copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114); 207 | } 208 | } 209 | else { 210 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 211 | } 212 | return dstimg; 213 | } 214 | 215 | void YOLOv5::nms(std::vector& input_boxes) 216 | { 217 | 218 | sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; }); // 降序排列 219 | std::vector remove_flags(input_boxes.size(),false); 220 | auto iou = [](const BoxInfo& box1,const BoxInfo& box2) 221 | { 222 | float xx1 = max(box1.x1, box2.x1); 223 | float yy1 = max(box1.y1, box2.y1); 224 | float xx2 = min(box1.x2, box2.x2); 225 | float yy2 = min(box1.y2, box2.y2); 226 | // 交集 227 | float w = max(0.0f, xx2 - xx1 + 1); 228 | float h = max(0.0f, yy2 - yy1 + 1); 229 | float inter_area = w * h; 230 | // 并集 231 | float union_area = max(0.0f,box1.x2-box1.x1) * max(0.0f,box1.y2-box1.y1) 232 | + max(0.0f,box2.x2-box2.x1) * max(0.0f,box2.y2-box2.y1) - inter_area; 233 | return inter_area / union_area; 234 | }; 235 | for (int i = 0; i < input_boxes.size(); ++i) 236 | { 237 | if(remove_flags[i]) continue; 238 | for (int j = i + 1; j < input_boxes.size(); ++j) 239 | { 240 | if(remove_flags[j]) continue; 241 | if(input_boxes[i].label == input_boxes[j].label && iou(input_boxes[i],input_boxes[j])>=this->nmsThreshold) 242 | { 243 | remove_flags[j] = true; 244 | } 245 | } 246 | } 247 | int idx_t = 0; 248 | // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素 249 | input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &remove_flags](const BoxInfo& f) { return remove_flags[idx_t++]; }), input_boxes.end()); 250 | } 251 | 252 | void YOLOv5::detect(Mat& frame) 253 | { 254 | int newh = 0, neww = 0, padh = 0, padw = 0; 255 | Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw); 256 | cv::cvtColor(dstimg, dstimg, cv::COLOR_BGR2RGB); // 由BGR转成RGB 257 | cv::Mat m_Normalized; 258 | dstimg.convertTo(m_Normalized, CV_32FC3, 1/255.); 259 | cv::split(m_Normalized, m_InputWrappers); // 通道分离[h,w,3] rgb 260 | auto ret = cudaMemcpyAsync(m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); 261 | auto ret1 = m_CudaContext->enqueueV2(m_ArrayDevMemory, m_CudaStream, nullptr); // TensorRT 执行通常是异步的，因此将内核排入 CUDA 流： 262 | ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], cudaMemcpyDeviceToHost, m_CudaStream); //输出传回给CPU，数据从显存到内存 263 | ret = cudaStreamSynchronize(m_CudaStream); 264 | float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex]; 265 | 266 | std::vector generate_boxes; // BoxInfo自定义的结构体 267 | float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww; 268 | for(int i = 0; i < m_iBoxNums; ++i) // 遍历所有的num_pre_boxes 269 | { 270 | int index = i * (m_iClassNums + 5); // prob[b*num_pred_boxes*(classes+5)] 271 | float obj_conf = pdata[index + 4]; // 置信度分数 272 | if (obj_conf > this->objThreshold) // 大于阈值 273 | { 274 | float* max_class_pos = std::max_element(pdata + index + 5, pdata + index + 5 + m_iClassNums); // 275 | (*max_class_pos) *= obj_conf; // 最大的类别分数*置信度 276 | if ((*max_class_pos) > this->confThreshold) // 再次筛选 277 | { 278 | //const int class_idx = classIdPoint.x; 279 | float cx = pdata[index]; //x 280 | float cy = pdata[index+1]; //y 281 | float w = pdata[index+2]; //w 282 | float h = pdata[index+3]; //h 283 | 284 | float xmin = (cx - padw - 0.5 * w)*ratiow; 285 | float ymin = (cy - padh - 0.5 * h)*ratioh; 286 | float xmax = (cx - padw + 0.5 * w)*ratiow; 287 | float ymax = (cy - padh + 0.5 * h)*ratioh; 288 | 289 | generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, (*max_class_pos), max_class_pos-(pdata + index + 5) }); 290 | } 291 | } 292 | } 293 | 294 | // Perform non maximum suppression to eliminate redundant overlapping boxes with 295 | // lower confidences 296 | nms(generate_boxes); 297 | for (size_t i = 0; i < generate_boxes.size(); ++i) 298 | { 299 | int xmin = int(generate_boxes[i].x1); 300 | int ymin = int(generate_boxes[i].y1); 301 | rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2); 302 | std::string label = format("%.2f", generate_boxes[i].score); 303 | label = this->classes[generate_boxes[i].label] + ":" + label; 304 | putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1); 305 | } 306 | } 307 | 308 | int main(int argc,char *argv[]) 309 | { 310 | clock_t startTime,endTime; //计算时间 311 | Configuration yolo_nets = { 0.3, 0.5, 0.3,"yolov5s.engine" }; 312 | YOLOv5 yolo_model(yolo_nets); 313 | std::string imgpath = "bus.jpg"; 314 | Mat srcimg = imread(imgpath); 315 | 316 | double timeStart = (double)getTickCount(); 317 | startTime = clock();//计时开始 318 | yolo_model.detect(srcimg); 319 | endTime = clock();//计时结束 320 | double nTime = ((double)getTickCount() - timeStart) / getTickFrequency(); 321 | std::cout << "clock_running time is:" <<(double)(endTime - startTime) / CLOCKS_PER_SEC << "s" << std::endl; 322 | std::cout << "The run time is:" << (double)clock() /CLOCKS_PER_SEC<< "s" << std::endl; 323 | std::cout << "getTickCount_running time :" << nTime << "sec\n" << std::endl; 324 | // static const string kWinName = "Deep learning object detection in ONNXRuntime"; 325 | // namedWindow(kWinName, WINDOW_NORMAL); 326 | // imshow(kWinName, srcimg); 327 | imwrite("restult_trt.jpg",srcimg); 328 | // waitKey(0); 329 | // destroyAllWindows(); 330 | return 0; 331 | } 332 | -------------------------------------------------------------------------------- /tensorrt_cmake/src/yolov5_trt.cpp: -------------------------------------------------------------------------------- 1 | // yzy 2 | // yolov5进行tensorrt部署的源文件 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "yolov5_trt.h" 8 | 9 | 10 | // 命名空间 11 | using namespace cv; 12 | using namespace nvinfer1; 13 | 14 | 15 | // Logger for TRT info/warning/errors, https://github.com/onnx/onnx-tensorrt/blob/main/onnx_trt_backend.cpp 16 | class TRT_Logger : public nvinfer1::ILogger 17 | { 18 | nvinfer1::ILogger::Severity _verbosity; 19 | std::ostream* _ostream; 20 | 21 | public: 22 | TRT_Logger(Severity verbosity = Severity::kWARNING, std::ostream& ostream = std::cout) 23 | : _verbosity(verbosity) 24 | , _ostream(&ostream) 25 | { 26 | } 27 | void log(Severity severity, const char* msg) noexcept override 28 | { 29 | if (severity <= _verbosity) 30 | { 31 | time_t rawtime = std::time(0); 32 | char buf[256]; 33 | strftime(&buf[0], 256, "%Y-%m-%d %H:%M:%S", std::gmtime(&rawtime)); 34 | const char* sevstr = (severity == Severity::kINTERNAL_ERROR ? " BUG" : severity == Severity::kERROR 35 | ? " ERROR" 36 | : severity == Severity::kWARNING ? "WARNING" : severity == Severity::kINFO ? " INFO" 37 | : "UNKNOWN"); 38 | (*_ostream) << "[" << buf << " " << sevstr << "] " << msg << std::endl; 39 | } 40 | } 41 | }; 42 | 43 | // 判断文件是否形成 44 | static bool ifFileExists(const char *FileName) 45 | { 46 | struct stat my_stat; 47 | return (stat(FileName, &my_stat) == 0); 48 | } 49 | 50 | 51 | // 加载onnx文件 52 | void YOLOv5::loadOnnx(const std::string strModelName) 53 | { 54 | TRT_Logger gLogger; // 日志 55 | //根据tensorrt pipeline 构建网络 56 | IBuilder* builder = createInferBuilder(gLogger); // 网络元数据,用于搭建网络入口 57 | builder->setMaxBatchSize(1); // batchsize 58 | const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); // 显式批处理 59 | INetworkDefinition* network = builder->createNetworkV2(explicitBatch); // 定义模型 60 | nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger); // 使用nvonnxparser 定义一个可用的onnx解析器 61 | parser->parseFromFile(strModelName.c_str(), static_cast(ILogger::Severity::kWARNING)); // 解析onnx 62 | // 使用builder对象构建engine 63 | IBuilderConfig* config = builder->createBuilderConfig(); // 64 | // 特别重要的属性是最大工作空间大小 65 | config->setMaxWorkspaceSize(1ULL << 30); // 分配内存空间 66 | m_CudaEngine = builder->buildEngineWithConfig(*network, *config); // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重 67 | 68 | std::string strTrtName = strModelName; 69 | size_t sep_pos = strTrtName.find_last_of("."); 70 | strTrtName = strTrtName.substr(0, sep_pos) + ".trt"; // 71 | IHostMemory *gieModelStream = m_CudaEngine->serialize(); // 将引擎序列化 72 | std::string serialize_str; // 73 | std::ofstream serialize_output_stream; 74 | serialize_str.resize(gieModelStream->size()); 75 | // memcpy内存拷贝函数 ，从源内存地址的起始位置开始拷贝若干个字节到目标内存地址中 76 | memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size()); 77 | serialize_output_stream.open(strTrtName.c_str()); 78 | serialize_output_stream<createExecutionContext(); //执行上下文用于执行推理 81 | // 使用一次，销毁parser，network, builder, and config 82 | parser->destroy(); 83 | network->destroy(); 84 | config->destroy(); 85 | builder->destroy(); 86 | } 87 | 88 | 89 | void YOLOv5::loadTrt(const std::string strName) 90 | { 91 | TRT_Logger gLogger; 92 | // 序列化引擎被保留并保存到文件中 93 | m_CudaRuntime = createInferRuntime(gLogger); 94 | std::ifstream fin(strName); 95 | std::string cached_engine = ""; 96 | while (fin.peek() != EOF) 97 | { 98 | std::stringstream buffer; 99 | buffer << fin.rdbuf(); 100 | cached_engine.append(buffer.str()); 101 | } 102 | fin.close(); 103 | m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr); // runtime对象反序列化 104 | m_CudaContext = m_CudaEngine->createExecutionContext(); //可以查询引擎获取有关网络的输入和输出的张量信息--维度/数据格式/数据类型 105 | m_CudaRuntime->destroy(); 106 | } 107 | 108 | 109 | // 初始化 110 | YOLOv5::YOLOv5(Configuration config) 111 | { 112 | confThreshold = config.confThreshold; 113 | nmsThreshold = config.nmsThreshold; 114 | objThreshold = config.objThreshold; 115 | inpHeight = 640; 116 | inpWidth = 640; 117 | 118 | std::string model_path = config.modelpath; // 模型权重路径 119 | // 加载模型 120 | std::string strTrtName = config.modelpath; // 加载模型权重 121 | size_t sep_pos = model_path.find_last_of("."); 122 | strTrtName = model_path.substr(0, sep_pos) + ".engine"; // ".trt" 123 | if(ifFileExists(strTrtName.c_str())) 124 | { 125 | loadTrt(strTrtName); 126 | } 127 | else 128 | { 129 | loadOnnx(config.modelpath); 130 | } 131 | 132 | // 利用加载的模型获取输入输出信息 133 | // 使用输入和输出blob名来获取输入和输出索引 134 | m_iInputIndex = m_CudaEngine->getBindingIndex("images"); // 输入索引 135 | m_iOutputIndex = m_CudaEngine->getBindingIndex("output"); // 输出 136 | Dims dims_i = m_CudaEngine->getBindingDimensions(m_iInputIndex); // 输入， 137 | int size1 = dims_i.d[0] * dims_i.d[1] * dims_i.d[2] * dims_i.d[3]; // 展平 138 | m_InputSize = cv::Size(dims_i.d[3], dims_i.d[2]); // 输入尺寸(W,H) 139 | Dims dims_o = m_CudaEngine->getBindingDimensions(m_iOutputIndex); // 输出，维度[0,1,2,3]NHWC 140 | int size2 = dims_o.d[0] * dims_o.d[1] * dims_o.d[2]; // 所有大小 141 | m_iClassNums = dims_o.d[2] - 5; // [,,classes+5] 142 | m_iBoxNums = dims_o.d[1]; // [b,num_pre_boxes,classes+5] 143 | 144 | // 分配内存大小 145 | cudaMalloc(&m_ArrayDevMemory[m_iInputIndex], size1 * sizeof(float)); 146 | m_ArrayHostMemory[m_iInputIndex] = malloc(size1 * sizeof(float)); 147 | m_ArraySize[m_iInputIndex] = size1 *sizeof(float); 148 | cudaMalloc(&m_ArrayDevMemory[m_iOutputIndex], size2 * sizeof(float)); 149 | m_ArrayHostMemory[m_iOutputIndex] = malloc( size2 * sizeof(float)); 150 | m_ArraySize[m_iOutputIndex] = size2 *sizeof(float); 151 | 152 | // bgr 153 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex]); 154 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + sizeof(float) * dims_i.d[2] * dims_i.d[3] ); 155 | m_InputWrappers.emplace_back(dims_i.d[2], dims_i.d[3], CV_32FC1, m_ArrayHostMemory[m_iInputIndex] + 2 * sizeof(float) * dims_i.d[2] * dims_i.d[3]); 156 | 157 | //创建CUDA流,推理时TensorRT执行通常是异步的，因此将内核排入CUDA流 158 | cudaStreamCreate(&m_CudaStream); // 只需初始化一次即可 159 | } 160 | 161 | void YOLOv5::UnInit() 162 | { 163 | 164 | for(auto &p: m_ArrayDevMemory) 165 | { 166 | cudaFree(p); 167 | p = nullptr; 168 | } 169 | for(auto &p: m_ArrayHostMemory) 170 | { 171 | free(p); 172 | p = nullptr; 173 | } 174 | cudaStreamDestroy(m_CudaStream); 175 | //m_CudaContext->destroy(); // 这个报错 176 | m_CudaEngine->destroy(); 177 | 178 | } 179 | 180 | YOLOv5::~YOLOv5() 181 | { 182 | UnInit(); 183 | } 184 | 185 | 186 | 187 | Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left) 188 | { 189 | int srch = srcimg.rows, srcw = srcimg.cols; 190 | *newh = this->inpHeight; 191 | *neww = this->inpWidth; 192 | Mat dstimg; 193 | if (this->keep_ratio && srch != srcw) { 194 | float hw_scale = (float)srch / srcw; 195 | if (hw_scale > 1) { 196 | *newh = this->inpHeight; 197 | *neww = int(this->inpWidth / hw_scale); 198 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 199 | *left = int((this->inpWidth - *neww) * 0.5); 200 | copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114); 201 | } 202 | else { 203 | *newh = (int)this->inpHeight * hw_scale; 204 | *neww = this->inpWidth; 205 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 206 | *top = (int)(this->inpHeight - *newh) * 0.5; 207 | copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114); 208 | } 209 | } 210 | else { 211 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 212 | } 213 | return dstimg; 214 | } 215 | 216 | void YOLOv5::nms(std::vector& input_boxes) 217 | { 218 | 219 | sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; }); // 降序排列 220 | std::vector remove_flags(input_boxes.size(),false); 221 | auto iou = [](const BoxInfo& box1,const BoxInfo& box2) 222 | { 223 | float xx1 = max(box1.x1, box2.x1); 224 | float yy1 = max(box1.y1, box2.y1); 225 | float xx2 = min(box1.x2, box2.x2); 226 | float yy2 = min(box1.y2, box2.y2); 227 | // 交集 228 | float w = max(0.0f, xx2 - xx1 + 1); 229 | float h = max(0.0f, yy2 - yy1 + 1); 230 | float inter_area = w * h; 231 | // 并集 232 | float union_area = max(0.0f,box1.x2-box1.x1) * max(0.0f,box1.y2-box1.y1) 233 | + max(0.0f,box2.x2-box2.x1) * max(0.0f,box2.y2-box2.y1) - inter_area; 234 | return inter_area / union_area; 235 | }; 236 | for (int i = 0; i < input_boxes.size(); ++i) 237 | { 238 | if(remove_flags[i]) continue; 239 | for (int j = i + 1; j < input_boxes.size(); ++j) 240 | { 241 | if(remove_flags[j]) continue; 242 | if(input_boxes[i].label == input_boxes[j].label && iou(input_boxes[i],input_boxes[j])>=this->nmsThreshold) 243 | { 244 | remove_flags[j] = true; 245 | } 246 | } 247 | } 248 | int idx_t = 0; 249 | // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素 250 | input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &remove_flags](const BoxInfo& f) { return remove_flags[idx_t++]; }), input_boxes.end()); 251 | } 252 | 253 | void YOLOv5::detect(Mat& frame) 254 | { 255 | int newh = 0, neww = 0, padh = 0, padw = 0; 256 | Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw); 257 | cv::cvtColor(dstimg, dstimg, cv::COLOR_BGR2RGB); // 由BGR转成RGB 258 | cv::Mat m_Normalized; 259 | dstimg.convertTo(m_Normalized, CV_32FC3, 1/255.); 260 | cv::split(m_Normalized, m_InputWrappers); // 通道分离[h,w,3] rgb 261 | auto ret = cudaMemcpyAsync(m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); 262 | auto ret1 = m_CudaContext->enqueueV2(m_ArrayDevMemory, m_CudaStream, nullptr); // TensorRT 执行通常是异步的，因此将内核排入 CUDA 流： 263 | ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], cudaMemcpyDeviceToHost, m_CudaStream); //输出传回给CPU，数据从显存到内存 264 | ret = cudaStreamSynchronize(m_CudaStream); 265 | float* pdata = (float*)m_ArrayHostMemory[m_iOutputIndex]; 266 | 267 | std::vector generate_boxes; // BoxInfo自定义的结构体 268 | float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww; 269 | for(int i = 0; i < m_iBoxNums; ++i) // 遍历所有的num_pre_boxes 270 | { 271 | int index = i * (m_iClassNums + 5); // prob[b*num_pred_boxes*(classes+5)] 272 | float obj_conf = pdata[index + 4]; // 置信度分数 273 | if (obj_conf > this->objThreshold) // 大于阈值 274 | { 275 | float* max_class_pos = std::max_element(pdata + index + 5, pdata + index + 5 + m_iClassNums); // 276 | (*max_class_pos) *= obj_conf; // 最大的类别分数*置信度 277 | if ((*max_class_pos) > this->confThreshold) // 再次筛选 278 | { 279 | //const int class_idx = classIdPoint.x; 280 | float cx = pdata[index]; //x 281 | float cy = pdata[index+1]; //y 282 | float w = pdata[index+2]; //w 283 | float h = pdata[index+3]; //h 284 | 285 | float xmin = (cx - padw - 0.5 * w)*ratiow; 286 | float ymin = (cy - padh - 0.5 * h)*ratioh; 287 | float xmax = (cx - padw + 0.5 * w)*ratiow; 288 | float ymax = (cy - padh + 0.5 * h)*ratioh; 289 | 290 | generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, (*max_class_pos), max_class_pos-(pdata + index + 5) }); 291 | } 292 | } 293 | } 294 | 295 | // Perform non maximum suppression to eliminate redundant overlapping boxes with 296 | // lower confidences 297 | nms(generate_boxes); 298 | for (size_t i = 0; i < generate_boxes.size(); ++i) 299 | { 300 | int xmin = int(generate_boxes[i].x1); 301 | int ymin = int(generate_boxes[i].y1); 302 | rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2); 303 | std::string label = format("%.2f", generate_boxes[i].score); 304 | label = this->classes[generate_boxes[i].label] + ":" + label; 305 | putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1); 306 | } 307 | } 308 | 309 | int main(int argc,char *argv[]) 310 | { 311 | clock_t startTime,endTime; //计算时间 312 | Configuration yolo_nets = { 0.3, 0.5, 0.3,"yolov5s.engine" }; 313 | YOLOv5 yolo_model(yolo_nets); 314 | std::string imgpath = "bus.jpg"; 315 | Mat srcimg = imread(imgpath); 316 | 317 | double timeStart = (double)getTickCount(); 318 | startTime = clock();//计时开始 319 | yolo_model.detect(srcimg); 320 | endTime = clock();//计时结束 321 | double nTime = ((double)getTickCount() - timeStart) / getTickFrequency(); 322 | std::cout << "clock_running time is:" <<(double)(endTime - startTime) / CLOCKS_PER_SEC << "s" << std::endl; 323 | std::cout << "The run time is:" << (double)clock() /CLOCKS_PER_SEC<< "s" << std::endl; 324 | std::cout << "getTickCount_running time :" << nTime << "sec\n" << std::endl; 325 | // static const string kWinName = "Deep learning object detection in ONNXRuntime"; 326 | // namedWindow(kWinName, WINDOW_NORMAL); 327 | // imshow(kWinName, srcimg); 328 | imwrite("restult_trt.jpg",srcimg); 329 | // waitKey(0); 330 | // destroyAllWindows(); 331 | return 0; 332 | } 333 | -------------------------------------------------------------------------------- /onnxruntime/ort.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include // 提供cuda加速 7 | #include // C或c++的api 8 | 9 | // 命名空间 10 | using namespace std; 11 | using namespace cv; 12 | using namespace Ort; 13 | 14 | // 自定义配置结构 15 | struct Configuration 16 | { 17 | public: 18 | float confThreshold; // Confidence threshold 19 | float nmsThreshold; // Non-maximum suppression threshold 20 | float objThreshold; //Object Confidence threshold 21 | string modelpath; 22 | }; 23 | 24 | // 定义BoxInfo结构类型 25 | typedef struct BoxInfo 26 | { 27 | float x1; 28 | float y1; 29 | float x2; 30 | float y2; 31 | float score; 32 | int label; 33 | } BoxInfo; 34 | 35 | // int endsWith(string s, string sub) { 36 | // return s.rfind(sub) == (s.length() - sub.length()) ? 1 : 0; 37 | // } 38 | 39 | // const float anchors_640[3][6] = { {10.0, 13.0, 16.0, 30.0, 33.0, 23.0}, 40 | // {30.0, 61.0, 62.0, 45.0, 59.0, 119.0}, 41 | // {116.0, 90.0, 156.0, 198.0, 373.0, 326.0} }; 42 | 43 | // const float anchors_1280[4][6] = { {19, 27, 44, 40, 38, 94},{96, 68, 86, 152, 180, 137},{140, 301, 303, 264, 238, 542}, 44 | // {436, 615, 739, 380, 925, 792} }; 45 | 46 | class YOLOv5 47 | { 48 | public: 49 | YOLOv5(Configuration config); 50 | void detect(Mat& frame); 51 | private: 52 | float confThreshold; 53 | float nmsThreshold; 54 | float objThreshold; 55 | int inpWidth; 56 | int inpHeight; 57 | int nout; 58 | int num_proposal; 59 | int num_classes; 60 | string classes[80] = {"person", "bicycle", "car", "motorbike", "aeroplane", "bus", 61 | "train", "truck", "boat", "traffic light", "fire hydrant", 62 | "stop sign", "parking meter", "bench", "bird", "cat", "dog", 63 | "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", 64 | "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", 65 | "skis", "snowboard", "sports ball", "kite", "baseball bat", 66 | "baseball glove", "skateboard", "surfboard", "tennis racket", 67 | "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", 68 | "banana", "apple", "sandwich", "orange", "broccoli", "carrot", 69 | "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", 70 | "bed", "diningtable", "toilet", "tvmonitor", "laptop", "mouse", 71 | "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", 72 | "sink", "refrigerator", "book", "clock", "vase", "scissors", 73 | "teddy bear", "hair drier", "toothbrush"}; 74 | 75 | const bool keep_ratio = true; 76 | vector input_image_; // 输入图片 77 | void normalize_(Mat img); // 归一化函数 78 | void nms(vector& input_boxes); 79 | Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left); 80 | 81 | Env env = Env(ORT_LOGGING_LEVEL_ERROR, "yolov5-6.1"); // 初始化环境 82 | Session *ort_session = nullptr; // 初始化Session指针选项 83 | SessionOptions sessionOptions = SessionOptions(); //初始化Session对象 84 | //SessionOptions sessionOptions; 85 | vector input_names; // 定义一个字符指针vector 86 | vector output_names; // 定义一个字符指针vector 87 | vector> input_node_dims; // >=1 outputs ，二维vector 88 | vector> output_node_dims; // >=1 outputs ,int64_t C/C++标准 89 | }; 90 | 91 | YOLOv5::YOLOv5(Configuration config) 92 | { 93 | this->confThreshold = config.confThreshold; 94 | this->nmsThreshold = config.nmsThreshold; 95 | this->objThreshold = config.objThreshold; 96 | this->num_classes = sizeof(this->classes)/sizeof(this->classes[0]); // 类别数量 97 | this->inpHeight = 640; 98 | this->inpWidth = 640; 99 | 100 | string model_path = config.modelpath; 101 | //std::wstring widestr = std::wstring(model_path.begin(), model_path.end()); //用于UTF-16编码的字符 102 | 103 | //gpu, https://blog.csdn.net/weixin_44684139/article/details/123504222 104 | //CUDA加速开启 105 | OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0); 106 | 107 | sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC); //设置图优化类型 108 | //ort_session = new Session(env, widestr.c_str(), sessionOptions); // 创建会话，把模型加载到内存中 109 | //ort_session = new Session(env, (const ORTCHAR_T*)model_path.c_str(), sessionOptions); // 创建会话，把模型加载到内存中 110 | ort_session = new Session(env, (const char*)model_path.c_str(), sessionOptions); 111 | size_t numInputNodes = ort_session->GetInputCount(); //输入输出节点数量 112 | size_t numOutputNodes = ort_session->GetOutputCount(); 113 | AllocatorWithDefaultOptions allocator; // 配置输入输出节点内存 114 | for (int i = 0; i < numInputNodes; i++) 115 | { 116 | input_names.push_back(ort_session->GetInputName(i, allocator)); // 内存 117 | Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i); // 类型 118 | auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo(); // 119 | auto input_dims = input_tensor_info.GetShape(); // 输入shape 120 | input_node_dims.push_back(input_dims); // 保存 121 | } 122 | for (int i = 0; i < numOutputNodes; i++) 123 | { 124 | output_names.push_back(ort_session->GetOutputName(i, allocator)); 125 | Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i); 126 | auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo(); 127 | auto output_dims = output_tensor_info.GetShape(); 128 | output_node_dims.push_back(output_dims); 129 | } 130 | this->inpHeight = input_node_dims[0][2]; 131 | this->inpWidth = input_node_dims[0][3]; 132 | this->nout = output_node_dims[0][2]; // 5+classes 133 | this->num_proposal = output_node_dims[0][1]; // pre_box 134 | 135 | } 136 | 137 | Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left) 138 | { 139 | int srch = srcimg.rows, srcw = srcimg.cols; 140 | *newh = this->inpHeight; 141 | *neww = this->inpWidth; 142 | Mat dstimg; 143 | if (this->keep_ratio && srch != srcw) { 144 | float hw_scale = (float)srch / srcw; 145 | if (hw_scale > 1) { 146 | *newh = this->inpHeight; 147 | *neww = int(this->inpWidth / hw_scale); 148 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 149 | *left = int((this->inpWidth - *neww) * 0.5); 150 | copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114); 151 | } 152 | else { 153 | *newh = (int)this->inpHeight * hw_scale; 154 | *neww = this->inpWidth; 155 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 156 | *top = (int)(this->inpHeight - *newh) * 0.5; 157 | copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114); 158 | } 159 | } 160 | else { 161 | resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA); 162 | } 163 | return dstimg; 164 | } 165 | 166 | void YOLOv5::normalize_(Mat img) 167 | { 168 | // img.convertTo(img, CV_32F); 169 | int row = img.rows; 170 | int col = img.cols; 171 | this->input_image_.resize(row * col * img.channels()); // vector大小 172 | for (int c = 0; c < 3; c++) // bgr 173 | { 174 | for (int i = 0; i < row; i++) // 行 175 | { 176 | for (int j = 0; j < col; j++) // 列 177 | { 178 | float pix = img.ptr(i)[j * 3 + 2 - c]; // Mat里的ptr函数访问任意一行像素的首地址,2-c:表示rgb 179 | this->input_image_[c * row * col + i * col + j] = pix / 255.0; 180 | 181 | } 182 | } 183 | } 184 | } 185 | 186 | void YOLOv5::nms(vector& input_boxes) 187 | { 188 | sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; }); // 降序排列 189 | vector vArea(input_boxes.size()); 190 | for (int i = 0; i < input_boxes.size(); ++i) 191 | { 192 | vArea[i] = (input_boxes[i].x2 - input_boxes[i].x1 + 1) 193 | * (input_boxes[i].y2 - input_boxes[i].y1 + 1); 194 | } 195 | // 全初始化为false，用来作为记录是否保留相应索引下pre_box的标志vector 196 | vector isSuppressed(input_boxes.size(), false); 197 | for (int i = 0; i < input_boxes.size(); ++i) 198 | { 199 | if (isSuppressed[i]) { continue; } 200 | for (int j = i + 1; j < input_boxes.size(); ++j) 201 | { 202 | if (isSuppressed[j]) { continue; } 203 | float xx1 = max(input_boxes[i].x1, input_boxes[j].x1); 204 | float yy1 = max(input_boxes[i].y1, input_boxes[j].y1); 205 | float xx2 = min(input_boxes[i].x2, input_boxes[j].x2); 206 | float yy2 = min(input_boxes[i].y2, input_boxes[j].y2); 207 | 208 | float w = max(0.0f, xx2 - xx1 + 1); 209 | float h = max(0.0f, yy2 - yy1 + 1); 210 | float inter = w * h; // 交集 211 | if(input_boxes[i].label == input_boxes[j].label) 212 | { 213 | float ovr = inter / (vArea[i] + vArea[j] - inter); // 计算iou 214 | if (ovr >= this->nmsThreshold) 215 | { 216 | isSuppressed[j] = true; 217 | } 218 | } 219 | } 220 | } 221 | // return post_nms; 222 | int idx_t = 0; 223 | // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素 224 | input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfo& f) { return isSuppressed[idx_t++]; }), input_boxes.end()); 225 | // 另一种写法 226 | // sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; }); // 降序排列 227 | // vector remove_flags(input_boxes.size(),false); 228 | // auto iou = [](const BoxInfo& box1,const BoxInfo& box2) 229 | // { 230 | // float xx1 = max(box1.x1, box2.x1); 231 | // float yy1 = max(box1.y1, box2.y1); 232 | // float xx2 = min(box1.x2, box2.x2); 233 | // float yy2 = min(box1.y2, box2.y2); 234 | // // 交集 235 | // float w = max(0.0f, xx2 - xx1 + 1); 236 | // float h = max(0.0f, yy2 - yy1 + 1); 237 | // float inter_area = w * h; 238 | // // 并集 239 | // float union_area = max(0.0f,box1.x2-box1.x1) * max(0.0f,box1.y2-box1.y1) 240 | // + max(0.0f,box2.x2-box2.x1) * max(0.0f,box2.y2-box2.y1) - inter_area; 241 | // return inter_area / union_area; 242 | // }; 243 | // for (int i = 0; i < input_boxes.size(); ++i) 244 | // { 245 | // if(remove_flags[i]) continue; 246 | // for (int j = i + 1; j < input_boxes.size(); ++j) 247 | // { 248 | // if(remove_flags[j]) continue; 249 | // if(input_boxes[i].label == input_boxes[j].label && iou(input_boxes[i],input_boxes[j])>=this->nmsThreshold) 250 | // { 251 | // remove_flags[j] = true; 252 | // } 253 | // } 254 | // } 255 | // int idx_t = 0; 256 | // // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素 257 | // input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &remove_flags](const BoxInfo& f) { return remove_flags[idx_t++]; }), input_boxes.end()); 258 | } 259 | 260 | void YOLOv5::detect(Mat& frame) 261 | { 262 | int newh = 0, neww = 0, padh = 0, padw = 0; 263 | Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw); 264 | this->normalize_(dstimg); 265 | // 定义一个输入矩阵，int64_t是下面作为输入参数时的类型 266 | array input_shape_{ 1, 3, this->inpHeight, this->inpWidth }; 267 | 268 | //创建输入tensor 269 | auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); 270 | Value input_tensor_ = Value::CreateTensor(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size()); 271 | 272 | // 开始推理 273 | vector ort_outputs = ort_session->Run(RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size()); // 开始推理 274 | /////generate proposals 275 | vector generate_boxes; // BoxInfo自定义的结构体 276 | float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww; 277 | float* pdata = ort_outputs[0].GetTensorMutableData(); // GetTensorMutableData 278 | for(int i = 0; i < num_proposal; ++i) // 遍历所有的num_pre_boxes 279 | { 280 | int index = i * nout; // prob[b*num_pred_boxes*(classes+5)] 281 | float obj_conf = pdata[index + 4]; // 置信度分数 282 | if (obj_conf > this->objThreshold) // 大于阈值 283 | { 284 | int class_idx = 0; 285 | float max_class_socre = 0; 286 | for (int k = 0; k < this->num_classes; ++k) 287 | { 288 | if (pdata[k + index + 5] > max_class_socre) 289 | { 290 | max_class_socre = pdata[k + index + 5]; 291 | class_idx = k; 292 | } 293 | } 294 | max_class_socre *= obj_conf; // 最大的类别分数*置信度 295 | if (max_class_socre > this->confThreshold) // 再次筛选 296 | { 297 | //const int class_idx = classIdPoint.x; 298 | float cx = pdata[index]; //x 299 | float cy = pdata[index+1]; //y 300 | float w = pdata[index+2]; //w 301 | float h = pdata[index+3]; //h 302 | 303 | float xmin = (cx - padw - 0.5 * w)*ratiow; 304 | float ymin = (cy - padh - 0.5 * h)*ratioh; 305 | float xmax = (cx - padw + 0.5 * w)*ratiow; 306 | float ymax = (cy - padh + 0.5 * h)*ratioh; 307 | 308 | generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, max_class_socre, class_idx }); 309 | } 310 | } 311 | } 312 | 313 | // Perform non maximum suppression to eliminate redundant overlapping boxes with 314 | // lower confidences 315 | nms(generate_boxes); 316 | for (size_t i = 0; i < generate_boxes.size(); ++i) 317 | { 318 | int xmin = int(generate_boxes[i].x1); 319 | int ymin = int(generate_boxes[i].y1); 320 | rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2); 321 | string label = format("%.2f", generate_boxes[i].score); 322 | label = this->classes[generate_boxes[i].label] + ":" + label; 323 | putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1); 324 | } 325 | } 326 | 327 | int main(int argc,char *argv[]) 328 | { 329 | clock_t startTime,endTime; //计算时间 330 | Configuration yolo_nets = { 0.3, 0.5, 0.3,"yolov5s.onnx" }; 331 | YOLOv5 yolo_model(yolo_nets); 332 | string imgpath = "bus.jpg"; 333 | Mat srcimg = imread(imgpath); 334 | 335 | double timeStart = (double)getTickCount(); 336 | startTime = clock();//计时开始 337 | yolo_model.detect(srcimg); 338 | endTime = clock();//计时结束 339 | double nTime = ((double)getTickCount() - timeStart) / getTickFrequency(); 340 | cout << "clock_running time is:" <<(double)(endTime - startTime) / CLOCKS_PER_SEC << "s" << endl; 341 | cout << "The run time is:" << (double)clock() /CLOCKS_PER_SEC<< "s" << endl; 342 | cout << "getTickCount_running time :" << nTime << "sec\n" << endl; 343 | // static const string kWinName = "Deep learning object detection in ONNXRuntime"; 344 | // namedWindow(kWinName, WINDOW_NORMAL); 345 | // imshow(kWinName, srcimg); 346 | imwrite("restult_ort.jpg",srcimg); 347 | // waitKey(0); 348 | // destroyAllWindows(); 349 | return 0; 350 | } 351 | --------------------------------------------------------------------------------