yolov5-6.0 tensorrt推理
CMakeLists.txt
cmake_minimum_required(VERSION 3.0) # CMake最低版本要求，低于2.6的构建过程会被终止set(CMAKE_CXX_STANDARD 14) #opencv4以上加此句
set(CMAKE_CXX_STANDARD 14)
project(cmake_test) # 定义工程名称
find_package(OpenCV REQUIRED)
# find_package(PCL REQUIRED)
find_package(CUDA REQUIRED)
find_package(TensorRT REQUIRED)
include_directories(
${PROJECT_SOURCE_DIR}/include
# ${PCL_INCLUDE_DIRS}
${CUDA_INCLUDE_DIRS}
${TensorRT_INCLUDE_DIRS}
${TensorRT_INCLUDE_DIRS}/../samples/common
)
add_executable(infer src/main.cpp)
target_link_libraries( infer
${OpenCV_LIBS}
# ${PCL_LIBRARIES}
${CUDA_LIBRARIES}
${TensorRT_LIBS}) # 链接库
infer.h
#include
using namespace std;
string names[] = {"xxxxx"}; // 修改自己的标签
struct BOX
{
float x;
float y;
float width;
float height;
};
struct Object
{
BOX box;
int label;
float confidence;
};
bool cmp(Object &obj1, Object &obj2){
return obj1.confidence > obj2.confidence;
}
vector<list<Object>> NMS(std::vector<Object> objs, float iou_thres = 0.45){
//第一步：将所有矩形框按照不同的类别标签分组，组内按照置信度高低得分进行排序；
list<Object> obj_l;
vector<list<Object>> NMS_List;
int a = 0;
for(int i = 0; i < 1; i++){
for(auto j : objs)
{
if(j.label == i){
obj_l.push_back(j);
obj_l.sort(cmp); //依据置信度升序排序
a = 1;
}
}
if(a == 1){
NMS_List.push_back(std::move(obj_l));
a = 0;
}
}
//第二步：计算IOU
float x1, y1, x1_w, y1_h,x2, y2, x2_w, y2_h;
float x_box, y_box, x_w_box, y_h_box, w_box, h_box;
float S1,S2,SBOX,res_iou;
int row = NMS_List.size(); //行数 列数：NMS_List[0].size()
int tmp;
for(int i = 0; i < row ; i++) //不同分类的循环
{
tmp = 0;
list<Object>::iterator it = NMS_List[i].begin();
while(it != --NMS_List[i].end()){
x1 = it->box.x;
y1 = it->box.y;
x1_w = x1 + it->box.width;
y1_h = y1 + it->box.height;
while(it != --NMS_List[i].end())
{
it++;
x2 = it->box.x;
y2 = it->box.y;
x2_w = x2 + it->box.width;
y2_h = y2 + it->box.height;
//交集左上角坐标x_box,y_box 框1-x1和框2-x2的最大值 框1-y1和框2-y2的最大值
x_box = std::max(x1, x2);
y_box = std::max(y1, y2);
//交集右下角坐标x_w_box,y_h_box 框1-x1_w和框2-x2_w的最小值 框1-y1_h和框2-y2_h的最小值
x_w_box = std::min(x1_w, x2_w);
y_h_box = std::min(y1_h, y2_h);
//交集框宽高
w_box = x_w_box - x_box;
h_box = y_h_box - y_box;
//无交集情况
if(w_box <= 0 || h_box <= 0)
{
it = NMS_List[i].erase(it);
if(it == NMS_List[i].end()){break;}
it--;
continue;
}
//有交集，计算IOU
S1 = (x1_w - x1) * (y1_h - y1);
S2 = (x2_w - x2) * (y2_h - y2);
SBOX = w_box * h_box;
if((res_iou = SBOX / (S1 + S2 - SBOX)) > iou_thres){
it = NMS_List[i].erase(it);
if(it == NMS_List[i].end()){break;}
it--;
}
}
it = NMS_List[i].begin();
if(it == --NMS_List[i].end()){break;}
tmp++;
for(int z = 0; z < tmp; z++){
it++;
if(it == --NMS_List[i].end()){break;}
}
}
}
return NMS_List;
}
main.cpp
#include 是 OpenCV 中用于访问图像像素的一种方法，使用 at 获取彩色图像中特定位置的像素颜色值
}
}
}
//二、模型反序列化
MyLogger logger;
//读取trt信息
const std::string engine_file_path = "../data/best.engine"; //填写自己trt文件路径
std::stringstream engine_file_stream;
engine_file_stream.seekg(0, engine_file_stream.beg); //从起始位置偏移0个字节，指针移动到文件流的开头
std::ifstream ifs(engine_file_path);
engine_file_stream << ifs.rdbuf();
ifs.close();
engine_file_stream.seekg(0, std::ios::end); //先把文件输入流指针定位到文档末尾来获取文档的长度
const int model_size = engine_file_stream.tellg(); //获取文件流的总长度
engine_file_stream.seekg(0, std::ios::beg);
void *model_mem = malloc(model_size); //开辟一样长的空间
engine_file_stream.read(static_cast<char *>(model_mem), model_size); //将内容读取到model_mem中
nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine *engine = runtime->deserializeCudaEngine(model_mem, model_size);
free(model_mem);
//三、模型推理
nvinfer1::IExecutionContext *context = engine->createExecutionContext();
void *buffers[2];
// 获取模型输入尺寸并分配GPU内存
nvinfer1::Dims input_dim = engine->getBindingDimensions(0);
int input_size = 1;
for (int j = 0; j < input_dim.nbDims; ++j) {
input_size *= input_dim.d[j];
}
cudaMalloc(&buffers[0], input_size * sizeof(float));
// 获取模型输出尺寸并分配GPU内存
nvinfer1::Dims output_dim = engine->getBindingDimensions(1);
int output_size = 1;
for (int j = 0; j < output_dim.nbDims; ++j) {
output_size *= output_dim.d[j];
}
cudaMalloc(&buffers[1], output_size * sizeof(float));
// 给模型输出数据分配相应的CPU内存
float *output_buffer = new float[output_size];
//
cudaStream_t stream;
cudaStreamCreate(&stream);
// 拷贝输入数据
cudaMemcpyAsync(buffers[0], input_blob,input_size * sizeof(float),
cudaMemcpyHostToDevice, stream);
// 投入数据流、执行推理
if(context->enqueueV2(buffers, stream, nullptr))
{
cout << "enqueueV2执行推理成功" << endl;
}
else{
cout << "enqueueV2执行推理失败" << endl;
return -1;
}
// 拷贝输出数据
cudaMemcpyAsync(output_buffer, buffers[1],output_size * sizeof(float),
cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);
delete context;
delete engine;
delete runtime;
delete[] input_blob;
//四、输出结果output_buffer，放入objs xywh为中心点坐标 和宽高
int nc = 1; // 自己的识别目标种类数
float *ptr = output_buffer;
std::vector<Object> objs;
for (int i = 0; i < 25200; ++i) {
const float objectness = ptr[4];
if (objectness >= 0.45f) {
const int label = std::max_element(ptr + 5, ptr + 5 + nc) - (ptr + 5); //std::max_element返回范围内的最大元素
const float confidence = ptr[5 + label] * objectness;
if (confidence >= 0.25f) {
const float bx = ptr[0];
const float by = ptr[1];
const float bw = ptr[2];
const float bh = ptr[3];
// std::cout << bx << "," << by << "," << bw << "," << bh << std::endl;
Object obj;
// 还原图像尺寸中box的尺寸比例，这里要减掉偏移值，并把box中心点坐标xy转成左上角坐标xy
obj.box.x = (bx - bw * 0.5f - x_offset) / ratio;
obj.box.y = (by - bh * 0.5f - y_offset) / ratio;
obj.box.width = bw / ratio;
obj.box.height = bh / ratio;
obj.label = label;
obj.confidence = confidence;
// std::cout << obj.box.x << "," << obj.box.y << "," << obj.box.width << "," << obj.box.height
// << "," << obj.label << "," << obj.confidence << std::endl;
objs.push_back(std::move(obj));
}
}
ptr += (5+nc);
} // i loop
//五、NMS非极大值抑制
vector<list<Object>> finalll = NMS(objs);
//六、画框
int row = finalll.size();
for(int i = 0; i < row; i++){
list<Object>::iterator it = finalll[i].begin();
while(it != finalll[i].end()){
cv::Point topLeft(it->box.x, it->box.y);
cv::Point bottomRight(it->box.x + it->box.width, it->box.y + it->box.height);
cv::rectangle(input_image, topLeft, bottomRight, cv::Scalar(0, 0, 255), 2);
std::stringstream buff;
buff.precision(2); //覆盖默认精度,置信度保留2位小数
buff.setf(std::ios::fixed);
buff << it->confidence;
string text =names[it->label] + " " + buff.str();
cv::putText(input_image, text, topLeft, 0, 1, cv::Scalar(0, 255, 0), 2);
it++;
}
}
cv::imwrite("pig.jpg", input_image);
return 0;
}