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 shape
this->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_NMS
cv::dnn::NMSBoxesBatched(bboxes, scores, labels, score_thres, iou_thres, indices);
#else
cv::dnn::NMSBoxes(bboxes, scores, score_thres, iou_thres, indices);
#endif
int 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.cpp
yolov8.cpp
common.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:0
cudaSetDevice(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.65f
std::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;
}
321
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
