[CV学习笔记] yolo&tensorrt多线程推理-第一部分

1、前言

之前分享了利用FastDet&tensorrt多线程推理的代码，本想着继续学习yolo&tensorrt多线程的代码，但是现在shouxieai直接开源的该项目，而且还包含yolov8实例分割的代码。因此本文主要是对项目代码进行梳理，一方面加深自己对多线程、cuda编程的理解,另一方面希望给有需要的同学提供帮助。本文主要对项目整体的框架进行说明，后续的文章将会对代码细节更加详细的注释。
yolo&tensorrt项目:https://github.com/shouxieai/infer
fastdet&ncnn 项目:https://blog.csdn.net/weixin_42108183/article/details/129331588

2、运行该项目

首先下载项目
CMakeLists.txt

cmake_minimum_required(VERSION 2.6)
# set (CMAKE_C_COMPILER "/usr/bin/gcc")
# set (CMAKE_CXX_COMPILER "/usr/bin/g++")
project(demo_trt)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_BUILD_TYPE Debug)
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace)
set(HAS_PYTHON OFF)
# 修改为本机地址
set(TENSORRT_DIR "/home/rex/TensorRT-8.4.0.6")
find_package(CUDA REQUIRED)
find_package(OpenCV)
include_directories(
    ${PROJECT_SOURCE_DIR}/src
    ${PROJECT_SOURCE_DIR}/src/trt
    ${PROJECT_SOURCE_DIR}/src/trt/common
    ${OpenCV_INCLUDE_DIRS}
    ${CUDA_TOOLKIT_ROOT_DIR}/include
    ${TENSORRT_DIR}/include
    ${CUDNN_DIR}/include
)
link_directories(
    ${TENSORRT_DIR}/lib
    ${CUDA_TOOLKIT_ROOT_DIR}/lib64
    ${CUDNN_DIR}/lib
)

set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O0 -Wfatal-errors -pthread -w -g")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O0 -Xcompiler -fPIC -g -w ${CUDA_GEN_CODE}")
file(GLOB_RECURSE cpp_srcs ${PROJECT_SOURCE_DIR}/src/*.cpp)
file(GLOB_RECURSE cuda_srcs ${PROJECT_SOURCE_DIR}/src/*.cu)
cuda_add_library(plugin_list SHARED ${cuda_srcs})
target_link_libraries(plugin_list nvinfer nvinfer_plugin)
target_link_libraries(plugin_list cuda cublas cudart cudnn)
target_link_libraries(plugin_list pthread)
target_link_libraries(plugin_list ${OpenCV_LIBS})
add_executable(demo_infer ${cpp_srcs})
target_link_libraries(demo_infer plugin_list)

# 生成engine
sh /workspace/build.sh
# 修改文件中的路径为绝对路径
mkdirt build && cd build
cmake ..
make -j
./../workspace/demo_infer

3、文件结构

src

cpm.hpp 基于生产者消费者的多线程推理
infer.cu 推理结果实现
infer.hpp 推理接口封装
yolo.cu 函数实现
yolo.hpp 任务类型、返回结果定义、推理函数封装
main.cpp 主函数

4、main.cpp

// 单batch推理
void single_inference()
{
    
       
    // 推荐写绝对路径
    cv::Mat image = cv::imread("/home/rex/Desktop/infer-lastest/workspace/inference/car.jpg");
    auto yolo = yolo::load("/home/rex/Desktop/infer-lastest/workspace/yolov8n-seg.b1.transd.engine", yolo::Type::V8Seg);
    if (yolo == nullptr)
        return;

    auto objs = yolo->forward(cvimg(image));
    int i = 0;
    cv::Mat image_mask = image.clone();
    int img_h = image_mask.rows;
    int img_w = image_mask.cols;

    std::vector<cv::Scalar> color;
    srand(time(0));
    for (int i = 0; i < 80; i++)
    {
    
    
        int b = rand() % 256;
        int g = rand() % 256;
        int r = rand() % 256;
        color.push_back(cv::Scalar(b, g, r));
    }
    for (auto &obj : objs)
    {
    
       
        // 实例分割
        if (obj.seg)
        {
    
    

            cv::Mat mask = cv::Mat(obj.seg->height, obj.seg->width, CV_8UC1, obj.seg->data);
            cv::resize(mask, mask, cv::Size((obj.right - obj.left), (obj.bottom - obj.top)), cv::INTER_NEAREST);
            inRange(mask, 127, 255, mask);
            cv::imwrite(
                cv::format("%d_mask.jpg", i), mask);
            i++;
            cv::Rect rect(obj.left, obj.top, (obj.right - obj.left), (obj.bottom - obj.top));
            cv::Mat c = image_mask(rect);
            c.setTo(color[obj.class_label], mask);
        }

        uint8_t b, g, r;
        tie(b, g, r) = yolo::random_color(obj.class_label);
        cv::rectangle(image, cv::Point(obj.left, obj.top),
                      cv::Point(obj.right, obj.bottom), cv::Scalar(b, g, r), 5);

        auto name = cocolabels[obj.class_label];
        auto caption = cv::format("%s %.2f", name, obj.confidence);
        int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
        cv::rectangle(image, cv::Point(obj.left - 3, obj.top - 33),
                      cv::Point(obj.left + width, obj.top), cv::Scalar(b, g, r),
                      -1);
        cv::putText(image, caption, cv::Point(obj.left, obj.top - 5), 0, 1,
                    cv::Scalar::all(0), 2, 16);
    }
    // 实例分割
    addWeighted(image, 0.5, image_mask, 0.5, 0, image);
    printf("Save result to Result.jpg, %d objects\n", objs.size());
    cv::imwrite("Result.jpg", image);
}

5、总结

本文将简单的将infer项目运行起来了，初步理解该项目的运行方式，后面的文章将会对代码细节进行更加详细的学习！