对图像做前处理,以便按照yolov5的engine输入格式输入网络

对图像做前处理,以便按照yolov5的engine输入格式输入网络

源码如下:

// preprocess.cpp  
// 对图像做前处理,以便按照engine输入格式输入网络
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#include
#include
#include 

using namespace std;

static const int INPUT_H = 6;
static const int INPUT_W = 6;


//加工图片变成拥有batch的输入, tensorrt输入数据的格式是一维(一维向量/张量)
void ProcessImage(cv::Mat image, float input_data[]) {
    //只处理一张图片,总之结果为一维[batch*3*INPUT_W*INPUT_H]
    //以下代码为投机取巧了

    cv::resize(image, image, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
    std::vector InputImage;

    InputImage.push_back(image);

    int ImgCount = InputImage.size();

    //float input_data[BatchSize * 3 * INPUT_H * INPUT_W];
    for (int b = 0; b < ImgCount; b++) {
        cv::Mat img = InputImage.at(b);
        cout<<"img  :  "<在这里插入图片描述
因为包含如下代码

cv::resize(image, image, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);

所以首先会先将图片转为6*6的尺寸,Mat类型结果如下:
对图像做前处理,以便按照yolov5的engine输入格式输入网络_第1张图片

1 cv::imread 读进来的图片是bgr格式
cv::Mat img = cv::imread(path);  
2 将Mat数据类型(img)转为一维数据(data)
void ProcessImage(cv::Mat image, float input_data[])

input_data若长度为3 * h * w,顺序[h * w个red, h * w个green, h * w个green], 因为for训环遍历图像元素的时候先给的uc_pixel[2]。

// float input_data[BatchSize * 3 * INPUT_H * INPUT_W]; //一维的数据
input_data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0;
input_data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
input_data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
3 存在缺点 ProcessImage()的过程会丢掉image的第一个像素

​ 结果如图所示:
对图像做前处理,以便按照yolov5的engine输入格式输入网络_第2张图片

4 所有代码
// preprocess.cpp  
// 对图像做前处理,以便按照engine输入格式输入网络
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#include
#include
#include 

using namespace std;

static const int INPUT_H = 6;
static const int INPUT_W = 6;


//加工图片变成拥有batch的输入, tensorrt输入数据的格式是一维(一维向量/张量)
void ProcessImage(cv::Mat image, float input_data[]) {
    //只处理一张图片,总之结果为一维[batch*3*INPUT_W*INPUT_H]
    //以下代码为投机取巧了

    cv::resize(image, image, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
    std::vector InputImage;

    InputImage.push_back(image);

    int ImgCount = InputImage.size();

    // float input_data[BatchSize * 3 * INPUT_H * INPUT_W];
    // input_data若长度为3*h*w,顺序[h*w个blue, h*w个green, h*w个red],
    // 其实应该是rgb顺序, 因为for训话遍历图像元素的时候先给的uc_pixel[2]。
    for (int b = 0; b < ImgCount; b++) {
        cv::Mat img = InputImage.at(b);
        cout<<"img2  :  \n"<
5 对应cmakelist.txt
cmake_minimum_required(VERSION 2.6)

project(preprocess)

add_definitions(-std=c++11)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

include_directories(${PROJECT_SOURCE_DIR}/include)
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
# cuda
include_directories(/usr/local/cuda-11.6/include)
link_directories(/usr/local/cuda-11.6/lib64)
# tensorrt
include_directories(/home/package/TensorRT-8.2.5.1/include/)
link_directories(/home/package/TensorRT-8.2.5.1/lib/)

include_directories(/home/package/TensorRT-8.2.5.1/samples/common/)
#link_directories(/home/package/TensorRT-8.2.5.1/lib/stubs/)

# opencv
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})
add_executable(yolo ${PROJECT_SOURCE_DIR}/preprocess.cpp)
target_link_libraries(yolo nvinfer)
target_link_libraries(yolo cudart)
target_link_libraries(yolo ${OpenCV_LIBS})
#如果onnx2engine则需要如下库
#target_link_libraries(yolo /home/mec/hlj/package/TensorRT-8.2.5.1/lib/stubs/libnvonnxparser.so)

add_definitions(-O2 -pthread)

你可能感兴趣的:(YOLO,c++,算法)