接到一个项目,需要用c++和单片机通信,还要使用yolo模型来做到目标检测的任务,但目前网上的各种博客并没有完整的流程教程,让我在部署过程费了不少劲,也踩了不少坑(甚至一度把ubuntu干黑屏)。于是想把训练及部署过程记录下来,并留给后来者方便使用。(博主使用的系统是ubuntu20.04)
1.Yolo简介
作为一个经典且实用的目标检测模型,yolo的性能强大已无需多言,现在(2023.4.1)yolo模型已经推出到yolov8,但是推理速度上yolov5还是够快且够用,而且对各种外部硬件的适配性强,如oak相机支持的模型还在yolov5-yolov6,
所以博主这里选择yolov5进行训练。
2.onnxruntime简介
onnxruntime是微软推出的一款推理框架,我们可以很方便的利用它运行一个onnx模型,而且它支持多种运行后端,包括CPU,GPU,TensorRT,DML等。onnxruntime可以说是对onnx模型最原生的支持了,而且onnxruntime也有在C++上部署使用的相关库,
所以我们选择onnxruntime作为我们的推理框架进行部署。
3.Yolov5模型训练及转换
怎么训练呢?
对于yolov5模型的训练,其实选择哪个都可以,博主这里模型使用的是
github.com/ultralytics…
训练过程参考
这个链接
这个博主对训练过程提供了非常详细的说明。按照这个博主的教程可以得到训练后的.pt文件,先别急着在models里export转换模型。如果在这里转换之后得到的output可能是这样的。
可以看到这个图片中有三个output,这样的结果是不能用的,我们需要的output是类似于[1,25200,7]这样的结果,并且这个结果必须在第一位上,通过观察可以知道25200是上图中三个输出结果中三个先验框大小和数量相乘得到的而我们训练使用的yolov5模型export的转换结果却不是这样的总和。
那怎么转换呢?
很简单,通过观察yolov5中yolo部分的代码发现
这个代码中的return是有问题的,返回结果不是我们想要的存在类似与[1,25200,7]的结果,而是详细细分的。这里我们最好不要改动yolo部分的代码,这可能会牵连其他部分从而产生错误。
我们直接使用
github.com/ultralytics…
这个模型来进行转换
首先先下载这个模型,用vscode打开后找到
export.py
在vscode的终端中输入
python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ...
注意yolov5s.pt处填自己训练结果的pt文件地址。然后我们使用netron查看得到的onnx模型
惊奇的发现我们得到了我们想要的output!(win!)(注意最后的85只是80个labels和5个参数,自己训练的模型有自己设置的labels所以数量可能不一样)
到这里我们已经完成了yolov5模型的训练和转换,请带好自己的随身物品(onnx模型),我们要赶往下一站了
4.利用cmake向C++部署该onnx模型
在这一步我们可能会遇到很多难题,但是不要着急,我来带着你一步步走
1.下载onnxruntime-gpu版本
博主这里使用的是onnxruntime-gpu-1.7
这里选择onnxruntime-linux-x64-gpu-1.7.0.tgz
下载后解压,然后放在主目录里,我们可以看到打开文件后有include和lib两个文件夹。这个我们之后要在cmakelist里链接。
2.下载vscode并创建cmake工程
这是博主的main函数代码,里面包含了onnxruntime的使用和推理结果的处理
#pragma comment(lib, "k4a.lib")
#include <fstream>
#include <sstream>
#include <iostream>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <k4a/k4a.h>
#include <k4arecord/record.h>
#include <k4arecord/playback.h>
#include <k4a/k4a.hpp>
#include <stdio.h>
#include <stdlib.h>
#include <opencv2/opencv.hpp>
#include <cstdlib>
#include <omp.h>
#include <opencv2/highgui/highgui_c.h>
#include "opencv2/imgproc/imgproc_c.h"
#include<opencv2/imgproc/types_c.h>
#include <onnxruntime_cxx_api.h>
#include <onnxruntime_c_api.h>
#include <cuda_provider_factory.h>
using namespace std;
using namespace cv;
using namespace Ort;
struct Configuration
public:
float confThreshold;
float nmsThreshold;
float objThreshold;
string modelpath;
typedef struct BoxInfo
float x1;
float y1;
float x2;
float y2;
float score;
int label;
} BoxInfo;
class YOLOv5
public:
YOLOv5(Configuration config);
void detect(Mat& frame);
private:
float confThreshold;
float nmsThreshold;
float objThreshold;
int inpWidth;
int inpHeight;
int nout;
int num_proposal;
int num_classes;
string classes[1] = {"tower"};
const bool keep_ratio = true;
vector<float> input_image_;
void normalize_(Mat img);
void nms(vector<BoxInfo>& input_boxes);
Mat resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left);
Env env = Env(ORT_LOGGING_LEVEL_ERROR, "yolov5-6.1");
Session *ort_session = nullptr;
SessionOptions sessionOptions = SessionOptions();
vector<char*> input_names;
vector<char*> output_names;
vector<vector<int64_t>> input_node_dims;
vector<vector<int64_t>> output_node_dims;
YOLOv5::YOLOv5(Configuration config)
this->confThreshold = config.confThreshold;
this->nmsThreshold = config.nmsThreshold;
this->objThreshold = config.objThreshold;
this->num_classes = 1;
this->inpHeight = 640;
this->inpWidth = 640;
string model_path = config.modelpath;
OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);
sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
ort_session = new Session(env, (const char*)model_path.c_str(), sessionOptions);
size_t numInputNodes = ort_session->GetInputCount();
size_t numOutputNodes = ort_session->GetOutputCount();
AllocatorWithDefaultOptions allocator;
for (int i = 0; i < numInputNodes; i++)
input_names.push_back(ort_session->GetInputName(i, allocator));
Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
auto input_dims = input_tensor_info.GetShape();
input_node_dims.push_back(input_dims);
for (int i = 0; i < numOutputNodes; i++)
output_names.push_back(ort_session->GetOutputName(i, allocator));
Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
auto output_dims = output_tensor_info.GetShape();
output_node_dims.push_back(output_dims);
this->inpHeight = input_node_dims[0][2];
this->inpWidth = input_node_dims[0][3];
this->nout = output_node_dims[0][2];
this->num_proposal = output_node_dims[0][1];
Mat YOLOv5::resize_image(Mat srcimg, int *newh, int *neww, int *top, int *left)
int srch = srcimg.rows, srcw = srcimg.cols;
*newh = this->inpHeight;
*neww = this->inpWidth;
Mat dstimg;
if (this->keep_ratio && srch != srcw) {
float hw_scale = (float)srch / srcw;
if (hw_scale > 1) {
*newh = this->inpHeight;
*neww = int(this->inpWidth / hw_scale);
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
*left = int((this->inpWidth - *neww) * 0.5);
copyMakeBorder(dstimg, dstimg, 0, 0, *left, this->inpWidth - *neww - *left, BORDER_CONSTANT, 114);
else {
*newh = (int)this->inpHeight * hw_scale;
*neww = this->inpWidth;
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
*top = (int)(this->inpHeight - *newh) * 0.5;
copyMakeBorder(dstimg, dstimg, *top, this->inpHeight - *newh - *top, 0, 0, BORDER_CONSTANT, 114);
else {
resize(srcimg, dstimg, Size(*neww, *newh), INTER_AREA);
return dstimg;
void YOLOv5::normalize_(Mat img)
int row = img.rows;
int col = img.cols;
this->input_image_.resize(row * col * img.channels());
for (int c = 0; c < 3; c++)
for (int i = 0; i < row; i++)
for (int j = 0; j < col; j++)
float pix = img.ptr<uchar>(i)[j * 3 + 2 - c];
this->input_image_[c * row * col + i * col + j] = pix / 255.0;
void YOLOv5::nms(vector<BoxInfo>& input_boxes)
sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; });
vector<bool> remove_flags(input_boxes.size(),false);
auto iou = [](const BoxInfo& box1,const BoxInfo& box2)
float xx1 = max(box1.x1, box2.x1);
float yy1 = max(box1.y1, box2.y1);
float xx2 = min(box1.x2, box2.x2);
float yy2 = min(box1.y2, box2.y2);
float w = max(0.0f, xx2 - xx1 + 1);
float h = max(0.0f, yy2 - yy1 + 1);
float inter_area = w * h;
float union_area = max(0.0f,box1.x2-box1.x1) * max(0.0f,box1.y2-box1.y1)
+ max(0.0f,box2.x2-box2.x1) * max(0.0f,box2.y2-box2.y1) - inter_area;
return inter_area / union_area;
for (int i = 0; i < input_boxes.size(); ++i)
if(remove_flags[i]) continue;
for (int j = i + 1; j < input_boxes.size(); ++j)
if(remove_flags[j]) continue;
if(input_boxes[i].label == input_boxes[j].label && iou(input_boxes[i],input_boxes[j])>=this->nmsThreshold)
remove_flags[j] = true;
int idx_t = 0;
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());
void YOLOv5::detect(Mat& frame)
int newh = 0, neww = 0, padh = 0, padw = 0;
Mat dstimg = this->resize_image(frame, &newh, &neww, &padh, &padw);
this->normalize_(dstimg);
array<int64_t, 4> input_shape_{ 1, 3, this->inpHeight, this->inpWidth };
这一行代码的作用是创建一个指向CPU内存的分配器信息对象(AllocatorInfo),用于在运行时分配和释放CPU内存。
它调用了CreateCpu函数并传递两个参数:OrtDeviceAllocator和OrtMemTypeCPU。
其中,OrtDeviceAllocator表示使用默认的设备分配器,OrtMemTypeCPU表示在CPU上分配内存。
通过这个对象,我们可以在运行时为张量分配内存,并且可以保证这些内存在计算完成后被正确地释放,避免内存泄漏的问题。
auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
vector<Value> ort_outputs = ort_session->Run(RunOptions{ nullptr }, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size());
vector<BoxInfo> generate_boxes;
float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww;
float* pdata = ort_outputs[0].GetTensorMutableData<float>();
for(int i = 0; i < num_proposal; ++i)
int index = i * nout;
float obj_conf = pdata[index + 4];
if (obj_conf > this->objThreshold)
int class_idx = 0;
float max_class_socre = 0;
for (int k = 0; k < this->num_classes; ++k)
if (pdata[k + index + 5] > max_class_socre)
max_class_socre = pdata[k + index + 5];
class_idx = k;
max_class_socre *= obj_conf;
if (max_class_socre > this->confThreshold)
float cx = pdata[index];
float cy = pdata[index+1];
float w = pdata[index+2];
float h = pdata[index+3];
float xmin = (cx - padw - 0.5 * w)*ratiow;
float ymin = (cy - padh - 0.5 * h)*ratioh;
float xmax = (cx - padw + 0.5 * w)*ratiow;
float ymax = (cy - padh + 0.5 * h)*ratioh;
generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, max_class_socre, class_idx });
nms(generate_boxes);
for (size_t i = 0; i < generate_boxes.size(); ++i)
int xmin = int(generate_boxes[i].x1);
int ymin = int(generate_boxes[i].y1);
rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2);
string label = format("%.2f", generate_boxes[i].score);
label = this->classes[generate_boxes[i].label] + ":" + label;
putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
int main(int argc,char *argv[])
找到并打开 Azure Kinect 设备
const uint32_t device_count = k4a::device::get_installed_count();
if (0 == device_count) {
cout << "Error: no K4A devices found. " << endl;
return -1;
else {
std::cout << "Found " << device_count << " connected devices. " << std::endl;
if (1 != device_count)
std::cout << "Error: more than one K4A devices found. " << std::endl;
return -1;
else
std::cout << "Done: found 1 K4A device. " << std::endl;
k4a::device device = k4a::device::open(K4A_DEVICE_DEFAULT);
std::cout << "Done: open device. " << std::endl;
检索并保存 Azure Kinect 图像数据
k4a_device_configuration_t config = K4A_DEVICE_CONFIG_INIT_DISABLE_ALL;
config.camera_fps = K4A_FRAMES_PER_SECOND_15;
config.color_format = K4A_IMAGE_FORMAT_COLOR_BGRA32;
config.color_resolution = K4A_COLOR_RESOLUTION_720P;
config.depth_mode = K4A_DEPTH_MODE_NFOV_UNBINNED;
config.synchronized_images_only = true;
device.start_cameras(&config);
k4a::image rgbImage;
Mat cv_rgbImage_no_alpha;
Mat cv_rgbImage_with_alpha;
k4a::capture capture;
Configuration yolo_nets = { 0.3, 0.5, 0.3,"/home/xue/picture/modle/best.onnx" };
YOLOv5 yolo_model(yolo_nets);
clock_t startTime,endTime;
while(true)
if (device.get_capture(&capture, std::chrono::milliseconds(0)))
if (device.get_capture(&capture)) {
rgbImage = capture.get_color_image();
cv_rgbImage_with_alpha = cv::Mat(rgbImage.get_height_pixels(), rgbImage.get_width_pixels(), CV_8UC4,
(void *)rgbImage.get_buffer());
cv::cvtColor(cv_rgbImage_with_alpha, cv_rgbImage_no_alpha, cv::COLOR_BGRA2BGR);
double timeStart = (double)getTickCount();
startTime = clock();
yolo_model.detect(cv_rgbImage_no_alpha);
endTime = clock();
double nTime = ((double)getTickCount() - timeStart) / getTickFrequency();
cout << "clock_running time is:" <<(double)(endTime - startTime) / CLOCKS_PER_SEC << "s" << endl;
cout << "The run time is:" << (double)clock() /CLOCKS_PER_SEC<< "s" << endl;
cout << "getTickCount_running time :" << nTime << "sec\n" << endl;
cv::namedWindow("color", CV_WINDOW_NORMAL);
cv::imshow("color",cv_rgbImage_no_alpha);
cv_rgbImage_with_alpha.release();
cv_rgbImage_no_alpha.release();
capture.reset();
if (cv::waitKey(1) == ' ')
{
std::cout << "----------------------------------" << std::endl;
std::cout << "------------- closed -------------" << std::endl;
std::cout << "----------------------------------" << std::endl;
break;
cv::destroyAllWindows();
device.close();
rgbImage.reset();
capture.reset();
return 0;
博主使用了azure相机,用opencv做到了实时采集图片并推理。所以代码里面包含了一些k4a库的函数,可以删掉换成自己的视频流或者图片。
关于代码有很多详细的注释,可以大概读一下,包括非极大值抑制等。
下面是cmakelist的代码
# cmake needs this line
cmake_minimum_required(VERSION 3.9.1)
# Define project name
project(yourproject)
# Kinect DK相机
find_package(k4a REQUIRED)
# Find OpenCV
find_package(OpenCV REQUIRED)
#find_package(Threads REQUIRED)
#find_package(OpenMP REQUIRED)
# Add includes
include_directories( ${CMAKE_CURRENT_LIST_DIR} ) # 包含当前目录下我的头文件
include_directories( ${OpenCV_INCLUDE_DIRS} )
include_directories(include)
# Enable C++11
set(CMAKE_CXX_FLAGS "-std=c++11")
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CUDA_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED TRUE)
set(ONNXRUNTIME_INCLUDE_DIRS your_include)
set(ONNXRUNTIME_LIBS your_libonnxruntime.so)
include_directories(${ONNXRUNTIME_INCLUDE_DIRS})
# Link your application with other libraries
add_executable(${PROJECT_NAME} "main.cpp")
target_include_directories(${PROJECT_NAME} PUBLIC ${ONNXRUNTIME_INCLUDE_DIRS})
target_link_libraries(onnxrun1 ${CUDA_LIBRARIES})
target_link_libraries(${PROJECT_NAME} k4a::k4a ${ONNXRUNTIME_LIBS} ${OpenCV_LIBS})
这个代码包括了opencv k4a 和onnxruntime
注意:不要使用find_package来添加onnxruntime,因为找不到,把your_include替换成你自己刚下好的onnxruntime-gpu里的include目录地址,把your_libonnxruntime.so替换成lib里的libonnxruntime.so
3.别忘了NVIDIA驱动以及CUDA和CUDNN的安装!
1.NVIDIA驱动
英伟达驱动官网可以在这里下载驱动
具体安装过程可以参考zhuanlan.zhihu.com/p/344925795 ,装完一定不要忘了sudo service lightdm start
开启图形化界面,nvidia版本或者图形化界面的问题都会导致ubuntu黑屏,装驱动一定要小心细心!(博主的血泪经历)
如果你之前装了NVIDIA的驱动,检查一下它的版本,ubuntu一般自带的那个驱动是不能用来跑这个的。
博主的驱动版本:nvidia-smi 515.76 (可以终端输入nvidia-smi查看,记住表格中的CUDA Version后的数字,等会下载cuda时版本号必须小于等于这个版本)
博主这里选择的版本是cuda11.2版本
参考blog.csdn.net/weixin_4215… 这个博文详细记录了CUDA和cudnn的安装过程,记得版本不要搞错,cudnn的版本也要对应cuda和nvidia来整。
如果你安装完了cuda和cudnn,就可以开始快乐的cmake了!
3.cmake
我们就不详细介绍cmake相关的内容了,按照上面的代码放进去后就可以直接cmake . make ./yourproject了。
然后就是快乐的结果展示时间辣!
可以看到每一帧的处理速度在0.06s左右,可以达到实时监测辣!
后续需要添加什么功能就可以自己在c++中改辣!不喜欢用python处理数据的人的福音
博主为了个通信的功能以及不想用python写代码开启了这段部署之旅,中间经历了onnxruntime安装后不会cmake链接,链接后发现缺东少西,甚至连nvidia的驱动都没装就开始调用cuda的gpu加速,装nvidia驱动甚至差点把ubuntu搞坏。这段血与泪我不想让第二个人经历了,所以写下这个博客来记录和分享。
之后我可能会做一些其他方面的博客,如果有需要的话就尽情留言吧!
我是一只咸鱼_,一个做视觉的小白,文章对你有帮助的话就点赞收藏吧! kisskiss!