【深度学习】C++ Tensorrt Yolov8 目标检测推理

C++ tensorrt对yolov8目标检测模型进行推理。
Windows版本下只需要修改common.hpp对文件的判断S_ISREG 和对文件夹的判断S_ISDIR即可，非核心代码，不调用删掉都可以。亲测可行。

模型导出

python 导出onnx

from ultralytics import YOLO# Load the YOLOv8 model
model = YOLO("best.pt")# # Export the model to ONNX format
model.export(format="onnx", dynamic=False, simplify=True, imgsz = (640,640), opset=12, half=False, int8=False)  # creates 'yolov8n.onnx'

tensorRT自带bin下的trtexec导出engine模型

## export trt/
(base) xiaoxin@xiaoxin:/usr/local/TensorRT-8.6.1.6/bin$ sudo ./trtexec --onnx=/home/xiaoxin/Documents/ultralytics-main/best.onnx --saveEngine=/home/xiaoxin/Documents/ultralytics-main/best.engine --workspace=1024 --fp16# Key	Value	Description
# format	'torchscript'	format to export to
# imgsz	640	image size as scalar or (h, w) list, i.e. (640, 480)
# keras	False	use Keras for TF SavedModel export
# optimize	False	TorchScript: optimize for mobile
# half	False	FP16 quantization
# int8	False	INT8 quantization
# dynamic	False	ONNX/TF/TensorRT: dynamic axes
# simplify	False	ONNX: simplify model
# opset	None	ONNX: opset version (optional, defaults to latest)
# workspace	4	TensorRT: workspace size (GB)
# nms	False	CoreML: add NMS

代码

yolov8.h

#ifndef YOLOV8_H
#define YOLOV8_H
#include "NvInferPlugin.h"
#include "common.hpp"
#include "fstream"
using namespace det;#define _PRINT true
// #define BATCHED_NMS
// #define assert(_Expression) ((void)0)class YOLOv8 {
public:explicit YOLOv8(const std::string& engine_file_path);~YOLOv8();void                 makePipe(bool warmup = true);void                 copyFromMat(const cv::Mat& image);void                 copyFromMat(const cv::Mat& image, cv::Size& size);void                 letterBox(const cv::Mat& image, cv::Mat& out, cv::Size& size);void                 infer();void                 postprocess(std::vector<Object>& objs,float                score_thres = 0.25f,float                iou_thres   = 0.65f,int                  topk        = 100,int                  num_labels  = 1);static void          draw_objects(const cv::Mat&                                image,cv::Mat&                                      res,const std::vector<Object>&                    objs,const std::vector<std::string>&               CLASS_NAMES,const std::vector<std::vector<unsigned int>>& COLORS);public:int                  num_bindings;int                  num_inputs  = 0;int                  num_outputs = 0;std::vector<Binding> input_bindings;std::vector<Binding> output_bindings;std::vector<void*>   host_ptrs;std::vector<void*>   device_ptrs;PreParam pparam;Parameter param;private:nvinfer1::ICudaEngine*       engine  = nullptr;nvinfer1::IRuntime*          runtime = nullptr;nvinfer1::IExecutionContext* context = nullptr;cudaStream_t                 stream  = nullptr;Logger                       gLogger{nvinfer1::ILogger::Severity::kERROR};
};
#endif  // YOLOV8_H

yolov8.cpp

#include "yolov8.h"// init engine model
YOLOv8::YOLOv8(const std::string& engine_file_path)
{// 1. make sure this file can be open by binary mode.std::ifstream file(engine_file_path, std::ios::binary);if(!file.good()){if(_PRINT){std::cout << "[ERROR] can not open file, please check up your engine file!" << std::endl;}return;}// 2. move pointer to the end.file.seekg(0, std::ios::end);// 3. get the location of current pointer.auto size = file.tellg();// 4. move pointer to start.file.seekg(0, std::ios::beg);char* trtModelStream = new char[size];assert(trtModelStream);file.read(trtModelStream, size);file.close();// 5. create runtime object deserialization///    important tip   ///// in order to use initLibNvInferPlugins, link to nvinfer_plugin.so or nvinfer_plugin.dll.// if you have some errors in this method, check up your .so or .dll files. you can put them in program directory.initLibNvInferPlugins(&this->gLogger, "");this->runtime = nvinfer1::createInferRuntime(this->gLogger);assert(this->runtime != nullptr);this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);assert(this->engine != nullptr);delete[] trtModelStream;// 6. create some space to store intermediate activation values.this->context = this->engine->createExecutionContext();assert(this->context != nullptr);cudaStreamCreate(&this->stream);// 7. get number of input tensor and output tensor.this->num_bindings = this->engine->getNbBindings();// 8. get binding dimensions, this process can support different dimensions.for (int i = 0; i < this->num_bindings; ++i) {Binding            binding;nvinfer1::Dims     dims;nvinfer1::DataType dtype = this->engine->getBindingDataType(i);std::string        name  = this->engine->getBindingName(i);binding.name             = name;binding.dsize            = type_to_size(dtype);bool IsInput = engine->bindingIsInput(i);if (IsInput) {this->num_inputs += 1;dims         = this->engine->getProfileDimensions(i, 0, nvinfer1::OptProfileSelector::kMAX);binding.size = get_size_by_dims(dims);binding.dims = dims;this->input_bindings.push_back(binding);// set max opt shapethis->context->setBindingDimensions(i, dims);}else {dims         = this->context->getBindingDimensions(i);binding.size = get_size_by_dims(dims);binding.dims = dims;this->output_bindings.push_back(binding);this->num_outputs += 1;}}
}YOLOv8::~YOLOv8()
{this->context->destroy();this->engine->destroy();this->runtime->destroy();cudaStreamDestroy(this->stream);for (auto& ptr : this->device_ptrs) {CHECK(cudaFree(ptr));}for (auto& ptr : this->host_ptrs) {CHECK(cudaFreeHost(ptr));}
}// warm up.
void YOLOv8::makePipe(bool warmup)
{for (auto& bindings : this->input_bindings) {void* d_ptr;CHECK(cudaMallocAsync(&d_ptr, bindings.size * bindings.dsize, this->stream));this->device_ptrs.push_back(d_ptr);}for (auto& bindings : this->output_bindings) {void * d_ptr, *h_ptr;size_t size = bindings.size * bindings.dsize;CHECK(cudaMallocAsync(&d_ptr, size, this->stream));CHECK(cudaHostAlloc(&h_ptr, size, 0));this->device_ptrs.push_back(d_ptr);this->host_ptrs.push_back(h_ptr);}if (warmup) {for (int i = 0; i < 5; i++) {for (auto& bindings : this->input_bindings) {size_t size  = bindings.size * bindings.dsize;void*  h_ptr = malloc(size);memset(h_ptr, 0, size);CHECK(cudaMemcpyAsync(this->device_ptrs[0], h_ptr, size, cudaMemcpyHostToDevice, this->stream));free(h_ptr);}this->infer();}if(_PRINT){printf("model warmup 5 times\n");}}
}void YOLOv8::letterBox(const cv::Mat& image, cv::Mat& out, cv::Size& size)
{const float inp_h  = size.height;const float inp_w  = size.width;float       height = image.rows;float       width  = image.cols;float r    = std::min(inp_h / height, inp_w / width);int   padw = std::round(width * r);int   padh = std::round(height * r);cv::Mat tmp;if ((int)width != padw || (int)height != padh) {cv::resize(image, tmp, cv::Size(padw, padh));}else {tmp = image.clone();}float dw = inp_w - padw;float dh = inp_h - padh;dw /= 2.0f;dh /= 2.0f;int top    = int(std::round(dh - 0.1f));int bottom = int(std::round(dh + 0.1f));int left   = int(std::round(dw - 0.1f));int right  = int(std::round(dw + 0.1f));cv::copyMakeBorder(tmp, tmp, top, bottom, left, right, cv::BORDER_CONSTANT, {114, 114, 114});cv::dnn::blobFromImage(tmp, out, 1 / 255.f, cv::Size(), cv::Scalar(0, 0, 0), true, false, CV_32F);this->pparam.ratio  = 1 / r;this->pparam.dw     = dw;this->pparam.dh     = dh;this->pparam.height = height;this->pparam.width  = width;
}void YOLOv8::copyFromMat(const cv::Mat& image)
{cv::Mat  nchw;auto&    in_binding = this->input_bindings[0];auto     width      = in_binding.dims.d[3];auto     height     = in_binding.dims.d[2];cv::Size size{width, height};this->letterBox(image, nchw, size);this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, height, width}});CHECK(cudaMemcpyAsync(this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
}void YOLOv8::copyFromMat(const cv::Mat& image, cv::Size& size)
{cv::Mat nchw;this->letterBox(image, nchw, size);this->context->setBindingDimensions(0, nvinfer1::Dims{4, {1, 3, size.height, size.width}});CHECK(cudaMemcpyAsync(this->device_ptrs[0], nchw.ptr<float>(), nchw.total() * nchw.elemSize(), cudaMemcpyHostToDevice, this->stream));
}void YOLOv8::infer()
{this->context->enqueueV2(this->device_ptrs.data(), this->stream, nullptr);for (int i = 0; i < this->num_outputs; i++) {size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;CHECK(cudaMemcpyAsync(this->host_ptrs[i], this->device_ptrs[i + this->num_inputs], osize, cudaMemcpyDeviceToHost, this->stream));}cudaStreamSynchronize(this->stream);
}void YOLOv8::postprocess(std::vector<Object>& objs, float score_thres, float iou_thres, int topk, int num_labels)
{if(param.setPam){score_thres = param.score_thres;iou_thres = param.iou_thres;topk = param.topk;num_labels = param.num_labels;}objs.clear();auto num_channels = this->output_bindings[0].dims.d[1];auto num_anchors  = this->output_bindings[0].dims.d[2];auto& dw     = this->pparam.dw;auto& dh     = this->pparam.dh;auto& width  = this->pparam.width;auto& height = this->pparam.height;auto& ratio  = this->pparam.ratio;std::vector<cv::Rect> bboxes;std::vector<float>    scores;std::vector<int>      labels;std::vector<int>      indices;cv::Mat output = cv::Mat(num_channels, num_anchors, CV_32F, static_cast<float*>(this->host_ptrs[0]));output         = output.t();for (int i = 0; i < num_anchors; i++) {auto  row_ptr    = output.row(i).ptr<float>();auto  bboxes_ptr = row_ptr;auto  scores_ptr = row_ptr + 4;auto  max_s_ptr  = std::max_element(scores_ptr, scores_ptr + num_labels);float score      = *max_s_ptr;if (score > score_thres) {float x = *bboxes_ptr++ - dw;float y = *bboxes_ptr++ - dh;float w = *bboxes_ptr++;float h = *bboxes_ptr;float x0 = clamp((x - 0.5f * w) * ratio, 0.f, width);float y0 = clamp((y - 0.5f * h) * ratio, 0.f, height);float x1 = clamp((x + 0.5f * w) * ratio, 0.f, width);float y1 = clamp((y + 0.5f * h) * ratio, 0.f, height);int              label = max_s_ptr - scores_ptr;cv::Rect_<float> bbox;bbox.x      = x0;bbox.y      = y0;bbox.width  = x1 - x0;bbox.height = y1 - y0;bboxes.push_back(bbox);labels.push_back(label);scores.push_back(score);}}#ifdef BATCHED_NMScv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
#elsecv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
#endifint cnt = 0;for (auto& i : indices) {if (cnt >= topk) {break;}Object obj;obj.rect  = bboxes[i];obj.prob  = scores[i];obj.label = labels[i];objs.push_back(obj);cnt += 1;}
}void YOLOv8::draw_objects(const cv::Mat&                                image,cv::Mat&                                      res,const std::vector<Object>&                    objs,const std::vector<std::string>&               CLASS_NAMES,const std::vector<std::vector<unsigned int>>& COLORS)
{res = image.clone();for (auto& obj : objs) {cv::Scalar color = cv::Scalar(COLORS[obj.label][0], COLORS[obj.label][1], COLORS[obj.label][2]);cv::rectangle(res, obj.rect, color, 2);char text[256];sprintf(text, "%s %.1f%%", CLASS_NAMES[obj.label].c_str(), obj.prob * 100);int      baseLine   = 0;int x = (int)obj.rect.x;int y = (int)obj.rect.y + 1;y > res.rows ? res.rows : y;/ you can choose whether you need a background for text. // cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine);// cv::rectangle(res, cv::Rect(x, y, label_size.width, label_size.height + baseLine), {0, 0, 255}, -1);cv::putText(res, text, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, {0, 0, 255}, 1);}
}

common.hpp

#ifndef COMMON_HPP
#define COMMON_HPP
#include "NvInfer.h"
#include "opencv2/opencv.hpp"
#include <sys/stat.h>
#include <unistd.h>#define CHECK(call)                                                                                                    \do {                                                                                                               \const cudaError_t error_code = call;                                                                           \if (error_code != cudaSuccess) {                                                                               \printf("CUDA Error:\n");                                                                                   \printf("    File:       %s\n", __FILE__);                                                                  \printf("    Line:       %d\n", __LINE__);                                                                  \printf("    Error code: %d\n", error_code);                                                                \printf("    Error text: %s\n", cudaGetErrorString(error_code));                                            \exit(1);                                                                                                   \}                                                                                                              \} while (0)class Logger: public nvinfer1::ILogger 
{
public:nvinfer1::ILogger::Severity reportableSeverity;explicit Logger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kINFO):reportableSeverity(severity){}void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override{if (severity > reportableSeverity) {return;}switch (severity) {case nvinfer1::ILogger::Severity::kINTERNAL_ERROR:std::cerr << "INTERNAL_ERROR: ";break;case nvinfer1::ILogger::Severity::kERROR:std::cerr << "ERROR: ";break;case nvinfer1::ILogger::Severity::kWARNING:std::cerr << "WARNING: ";break;case nvinfer1::ILogger::Severity::kINFO:std::cerr << "INFO: ";break;default:std::cerr << "VERBOSE: ";break;}std::cerr << msg << std::endl;}
};inline int get_size_by_dims(const nvinfer1::Dims& dims)
{int size = 1;for (int i = 0; i < dims.nbDims; i++) {size *= dims.d[i];}return size;
}inline int type_to_size(const nvinfer1::DataType& dataType)
{switch (dataType) {case nvinfer1::DataType::kFLOAT:return 4;case nvinfer1::DataType::kHALF:return 2;case nvinfer1::DataType::kINT32:return 4;case nvinfer1::DataType::kINT8:return 1;case nvinfer1::DataType::kBOOL:return 1;default:return 4;}
}inline static float clamp(float val, float min, float max)
{return val > min ? (val < max ? val : max) : min;
}inline bool IsPathExist(const std::string& path)
{if (access(path.c_str(), 0) == F_OK) {return true;}return false;
}inline bool IsFile(const std::string& path)
{if (!IsPathExist(path)) {printf("%s:%d %s not exist\n", __FILE__, __LINE__, path.c_str());return false;}struct stat buffer;return (stat(path.c_str(), &buffer) == 0 && S_ISREG(buffer.st_mode));
}inline bool IsFolder(const std::string& path)
{if (!IsPathExist(path)) {return false;}struct stat buffer;return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
}namespace det 
{struct Binding {size_t         size  = 1;size_t         dsize = 1;nvinfer1::Dims dims;std::string    name;};struct Object {cv::Rect_<float> rect;int              label = 0;float            prob  = 0.0;};struct PreParam {float ratio  = 1.0f;float dw     = 0.0f;float dh     = 0.0f;float height = 0;float width  = 0;};struct Parameter{bool setPam = false;float score_thres = 0.25f;float iou_thres = 0.65f;int topk = 100;int num_labels = 1;};
}  // namespace det
#endif  // COMMON_HPP

CMakeList

cmake_minimum_required(VERSION 3.1)set(CMAKE_CUDA_ARCHITECTURES 60 61 62 70 72 75 86 89 90)
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)project(yolov8 LANGUAGES CXX CUDA)set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O3")
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_BUILD_TYPE Release)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)# CUDA
find_package(CUDA REQUIRED)
message(STATUS "CUDA Libs: \n${CUDA_LIBRARIES}\n")
get_filename_component(CUDA_LIB_DIR ${CUDA_LIBRARIES} DIRECTORY)
message(STATUS "CUDA Headers: \n${CUDA_INCLUDE_DIRS}\n")# OpenCV
find_package(OpenCV REQUIRED)
message(STATUS "OpenCV Libs: \n${OpenCV_LIBS}\n")
message(STATUS "OpenCV Libraries: \n${OpenCV_LIBRARIES}\n")
message(STATUS "OpenCV Headers: \n${OpenCV_INCLUDE_DIRS}\n")# TensorRT
set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu)
set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu)message(STATUS "TensorRT Libs: \n${TensorRT_LIBRARIES}\n")
message(STATUS "TensorRT Headers: \n${TensorRT_INCLUDE_DIRS}\n")list(APPEND INCLUDE_DIRS${CUDA_INCLUDE_DIRS}${OpenCV_INCLUDE_DIRS}${TensorRT_INCLUDE_DIRS}include)list(APPEND ALL_LIBS${CUDA_LIBRARIES}${CUDA_LIB_DIR}${OpenCV_LIBRARIES}${TensorRT_LIBRARIES})include_directories(${INCLUDE_DIRS})add_executable(${PROJECT_NAME}main.cppyolov8.cppcommon.hpp)target_link_directories(${PROJECT_NAME} PUBLIC ${ALL_LIBS})
target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer nvinfer_plugin cudart ${OpenCV_LIBS})if (${OpenCV_VERSION} VERSION_GREATER_EQUAL 4.7.0)message(STATUS "Build with -DBATCHED_NMS")add_definitions(-DBATCHED_NMS)
endif ()

main.cpp

#include "chrono"
#include "opencv2/opencv.hpp"
#include "yolov8.h"
#include <iostream>using namespace std;const std::vector<std::string> CLASS_NAMES = {"blackPoint"};const std::vector<std::vector<unsigned int>> COLORS = {{0, 0, 255}};int main(int argc, char** argv)
{// cuda:0cudaSetDevice(0);const std::string engine_file_path{"/home/xiaoxin/Documents/ultralytics-main/last.engine"};const std::string path{"/home/xiaoxin/Documents/ultralytics-main/datasets/Tray/labelImg"};std::vector<std::string> imagePathList;bool                     isVideo{false};auto yolov8 = new YOLOv8(engine_file_path);yolov8->makePipe(true);if (IsFile(path)){std::string suffix = path.substr(path.find_last_of('.') + 1);if (suffix == "jpg" || suffix == "jpeg" || suffix == "png") {imagePathList.push_back(path);}else if (suffix == "mp4" || suffix == "avi" || suffix == "m4v" || suffix == "mpeg" || suffix == "mov"|| suffix == "mkv") {isVideo = true;}else {printf("suffix %s is wrong !!!\n", suffix.c_str());std::abort();}}else if (IsFolder(path)) {cv::glob(path + "/*.png", imagePathList);}cv::Mat  res, image;cv::Size size        = cv::Size{640, 640};yolov8->param.setPam = true;yolov8->param.num_labels  = 1;yolov8->param.topk        = 100;yolov8->param.score_thres = 0.25f;yolov8->param.iou_thres   = 0.35f; // 0.65fstd::vector<Object> objs;cv::namedWindow("result", cv::WINDOW_AUTOSIZE);if (isVideo) {cv::VideoCapture cap(path);if (!cap.isOpened()) {printf("can not open %s\n", path.c_str());return -1;}while (cap.read(image)) {objs.clear();yolov8->copyFromMat(image, size);auto start = std::chrono::system_clock::now();yolov8->infer();auto end = std::chrono::system_clock::now();yolov8->postprocess(objs);yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;printf("cost %2.4lf ms\n", tc);cv::imshow("result", res);if (cv::waitKey(10) == 'q') {break;}}}else {for (auto& path : imagePathList) {objs.clear();image = cv::imread(path);yolov8->copyFromMat(image, size);auto start = std::chrono::system_clock::now();yolov8->infer();yolov8->postprocess(objs);yolov8->draw_objects(image, res, objs, CLASS_NAMES, COLORS);auto end = std::chrono::system_clock::now();auto tc = (double)std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.;printf("cost %2.4lf ms\n", tc);resize(res, res, cv::Size(0,0), 3, 3);cv::imshow("result", res);cv::waitKey(0);}}cv::destroyAllWindows();delete yolov8;return 0;
}