对图像做前处理，以便按照yolov5的engine输入格式输入网络
Posted 爱吃油淋鸡的莫何
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了对图像做前处理，以便按照yolov5的engine输入格式输入网络相关的知识，希望对你有一定的参考价值。
对图像做前处理，以便按照yolov5的engine输入格式输入网络
源码如下：
// preprocess.cpp
// 对图像做前处理，以便按照engine输入格式输入网络
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <cmath>
#include <cassert>
#include<opencv2/core/core.hpp>
#include<opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
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<cv::Mat> 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 : "<<img.size()<<endl;
int w = img.cols;
int h = img.rows;
int i = 0;
//对img中的像素值做归一化，将归一化后的像素保存在input_data中（input相当于一个一维向量）
for (int row = 0; row < h; ++row)
uchar* uc_pixel = img.data + row * img.step;
for (int col = 0; col < INPUT_W; ++col)
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;
uc_pixel += 3;
i++;
int main(int argc, char** argv)
// cout<<"argv = "<<argv<<endl;
// cout<<"argv[1] = "<<argv[1]<<endl;
std::string path = "/home/ccc/onnx2engine/yolov5_5.0/TinyImg.jpg";
std::cout << "img_path=" << path << endl;
static float data[3 * INPUT_H * INPUT_W];
cv::Mat img = cv::imread(path);
// 将Mat数据类型（img）转为一维数据(data)
ProcessImage(img, data);
return 0;
详细解释：
假设我输入的图像TinyImg.jpg的尺寸大小为6*8，打印读入的Mat img = cv::imread(path)，结果如下：
因为包含如下代码
cv::resize(image, image, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
所以首先会先将图片转为6*6的尺寸，Mat类型结果如下：
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的第一个像素
结果如图所示：
4 所有代码
// preprocess.cpp
// 对图像做前处理，以便按照engine输入格式输入网络
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include <cmath>
#include <cassert>
#include<opencv2/core/core.hpp>
#include<opencv2/highgui/highgui.hpp>
#include <opencv2/opencv.hpp>
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<cv::Mat> 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"<<img<<endl;
//cout<<"img2 : "<<img.size()<<endl;
int w = img.cols;
int h = img.rows;
int i = 0;
//对img中的像素值做归一化，将归一化后的像素保存在input_data中（input相当于一个一维向量）
// cout<<"***********************************"<<endl;
// cout<<"b w h i = "<<b<<" "<<w<<" "<<h<<" "<<i<<endl;
for (int row = 0; row < h; ++row)
// img.data 认为图像起始位置的指针 img.step = 每行的像素数 * 3
uchar* uc_pixel = img.data + row * img.step;
for (int col = 0; col < INPUT_W; ++col)
// cout<<" **********row col ="<<row<<" "<<col<<endl;
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;
uc_pixel += 3;
i++;
// cout <<" uc_pixel[2] = "<<(float)uc_pixel[2]<<endl;
// cout <<" uc_pixel[0] = "<<(float)uc_pixel[0]<<endl;
// cout <<" uc_pixel[1] = "<<(float)uc_pixel[1]<<endl;
// cout <<" input_data = "<<input_data<<endl;
int main(int argc, char** argv)
// cout<<"argv = "<<argv<<endl;
// cout<<"argv[1] = "<<argv[1]<<endl;
// std::string path = "/home/hlj/sources/data220729_0467.jpg";
std::string path = "/home/hlj/onnx2engine/yolov5_5.0/TinyImg.jpg";
std::cout << "img_path=" << path << endl;
static float data[3 * INPUT_H * INPUT_W];
cv::Mat img = cv::imread(path); // cv::imread 读进来的图片是bgr格式
cout<<"img1 : \\n"<<img<<endl;
// 将Mat数据类型（img）转为一维数据(data)
ProcessImage(img, data);
return 0;
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)
以上是关于对图像做前处理，以便按照yolov5的engine输入格式输入网络的主要内容，如果未能解决你的问题，请参考以下文章