yolov5 进行tensorrt加速 C++版

修改模型、标签以及测试图片的路径

main.cpp

#include 
#include "src/yolov5.h"

int main()
{
    std::vector<YOLOv5::DetectRes> result;
    std::string onnx_file = "/media/yao/Data/yolov5-tensorrt/models/yolov5s.onnx";
    std::string label_file = "/media/yao/Data/yolov5-tensorrt/src/coco.names";
    cv::Mat org_img = cv::imread("/media/yao/Data/yolov5-tensorrt/images/bus.jpg");
    YOLOv5 YOLOv5(onnx_file,label_file);
    YOLOv5.Init_Model();
    result = YOLOv5.Inference(org_img);
    for(const auto &rect : result)
    {
        std::string name = rect.classes;
        cv::putText(org_img, name, cv::Point(rect.x - rect.w / 2, rect.y - rect.h / 2 - 5), cv::FONT_HERSHEY_COMPLEX, 0.7, cv::Scalar(255,255,0), 2);
        cv::Rect rst(rect.x - rect.w / 2, rect.y - rect.h / 2, rect.w, rect.h);
        cv::rectangle(org_img, rst,cv::Scalar(255,255,0), 2, cv::LINE_8, 0);
    }
    cv::imwrite("1.jpg", org_img);
    return 0;
}

修改CMakeLists.txt中的TensorRT路径

cmake_minimum_required(VERSION 3.5)

project(yolov5_trt)

set(CMAKE_CXX_STANDARD 14)

# CUDA
find_package(CUDA REQUIRED)
message(STATUS "Find CUDA include at ${CUDA_INCLUDE_DIRS}")
message(STATUS "Find CUDA libraries: ${CUDA_LIBRARIES}")

# TensorRT
set(TENSORRT_ROOT /home/yao/TensorRT-7.1.3.4)
find_path(TENSORRT_INCLUDE_DIR NvInfer.h
        HINTS ${TENSORRT_ROOT} PATH_SUFFIXES include/)
message(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
find_library(TENSORRT_LIBRARY_INFER nvinfer
        HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
find_library(TENSORRT_LIBRARY_ONNXPARSER nvonnxparser
        HINTS  ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
        PATH_SUFFIXES lib lib64 lib/x64)
set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_ONNXPARSER})
message(STATUS "Find TensorRT libs: ${TENSORRT_LIBRARY}")

# OpenCV
find_package(OpenCV REQUIRED)
message(STATUS "Find OpenCV include at ${OpenCV_INCLUDE_DIRS}")
message(STATUS "Find OpenCV libraries: ${OpenCV_LIBRARIES}")

set(COMMON_INCLUDE ./common)

include_directories(${CUDA_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS} ${COMMON_INCLUDE})

add_executable(yolov5_trt main.cpp src/yolov5.cpp)
target_link_libraries(yolov5_trt ${OpenCV_LIBRARIES} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY})

演示视频：https://www.bilibili.com/video/BV1or4y177gS?spm_id_from=333.999.0.0&vd_source=4cd5ac8dda02d0b3152cd9b05f7e4006
完整代码：https://github.com/yaoyi30/yolov5_TensorRT_C-