一、说明
1、使用上一篇文章中生成模型的.engine文件,具体操作见上一篇博客:YOLOV8部署在jetson中--onnx转换为.engine文件
2、完整代码地址:YOLOV8-jetson完整代码
二、核心代码解析
推理的过程主要分为三部分:
1、engine文件读取、调用推理接口,得到context指针,后续推理使用。
if (argc < 2) {
argv[1] = "model_engine/yolov8n1_fp16.engine";
}
// create a model using the API directly and serialize it to a stream
char* trtModelStream{ nullptr }; //char* trtModelStream==nullptr; 开辟空指针后 要和new配合使用
size_t size{ 0 };//与int固定四个字节不同有所不同,size_t的取值range是目标平台下最大可能的数组尺寸,一些平台下size_t的范围小于int的正数范围,又或者大于unsigned int. 使用Int既有可能浪费,又有可能范围不够大。
std::ifstream file(argv[1], std::ios::binary);
if (file.good()) {
std::cout << "load engine success" << std::endl;
file.seekg(0, file.end);//指向文件的最后地址
size = file.tellg();//把文件长度告诉给size
//std::cout << "\nfile:" << argv[1] << " size is";
//std::cout << size << "";
file.seekg(0, file.beg);//指回文件的开始地址
trtModelStream = new char[size];//开辟一个char 长度是文件的长度
assert(trtModelStream);//
file.read(trtModelStream, size);//将文件内容传给trtModelStream
file.close();//关闭
}
else {
std::cout << "load engine failed" << std::endl;
return 1;
}
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
bool didInitPlugins = initLibNvInferPlugins(nullptr, "");
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
2、USB摄像头调用、处理每一帧图像数据。主要是对图像进行resize、pading、归一化操作。最后将3*640*640的图像放到数组data中。
VideoCapture cap(1); //终端输入:ls /dev/video* 来获取摄像头id
if(!cap.isOpened()){
std::cout<<"creamer error!"<<std::endl;
return -1;
}
double fps;
while(true){
Mat src;
cap.read(src);
// Mat src = imread(argv[2], 1);
// if (src.empty()) { std::cout << "image load faild" << std::endl; return 1; }
int img_width = src.cols;
int img_height = src.rows;
//std::cout << "宽高:" << img_width << " " << img_height << std::endl;
// Subtract mean from image
static float data[3 * INPUT_H * INPUT_W];
Mat pr_img0, pr_img;
std::vector<int> padsize;
pr_img = preprocess_img(src, INPUT_H, INPUT_W, padsize); // Resize
int newh = padsize[0], neww = padsize[1], padh = padsize[2], padw = padsize[3];
float ratio_h = (float)src.rows / newh;
float ratio_w = (float)src.cols / neww;
int i = 0;// [1,3,INPUT_H,INPUT_W]
//std::cout << "pr_img.step" << pr_img.step << std::endl;
for (int row = 0; row < INPUT_H; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;//pr_img.step=widthx3 就是每一行有width个3通道的值
for (int col = 0; col < INPUT_W; ++col)
{
data[i] = (float)uc_pixel[2] / 255.0;
data[i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
data[i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
uc_pixel += 3;
++i;
}
}
3、推理,将context、data传入到doinference推理,prob是用于保存推理结果的数组。
static float prob[OUTPUT_SIZE];
auto start = std::chrono::system_clock::now();
auto start0 = std::chrono::system_clock::now();
doInference(*context, data, prob, 1);
auto end = std::chrono::system_clock::now();
std::cout << "推理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
4、后处理,prob数组大小为705600,对应v8目标检测输出1*84*8400,根据80个类别的置信度得分画框。在该代码里,标签为0-79,0就是对应person。后处理细节部分可以自行查阅官方资料。
int net_length = CLASSES + 4;
cv::Mat out1 = cv::Mat(net_length, Num_box, CV_32F, prob);
start = std::chrono::system_clock::now();
for (int i = 0; i < Num_box; i++) {
//输出是1*net_length*Num_box;所以每个box的属性是每隔Num_box取一个值,共net_length个值
cv::Mat scores = out1(Rect(i, 4, 1, CLASSES)).clone();
Point classIdPoint;
double max_class_socre;
minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
max_class_socre = (float)max_class_socre;
if (max_class_socre >= CONF_THRESHOLD) {
float x = (out1.at<float>(0, i) - padw) * ratio_w; //cx
float y = (out1.at<float>(1, i) - padh) * ratio_h; //cy
float w = out1.at<float>(2, i) * ratio_w; //w
float h = out1.at<float>(3, i) * ratio_h; //h
int left = MAX((x - 0.5 * w), 0);
int top = MAX((y - 0.5 * h), 0);
int width = (int)w;
int height = (int)h;
if (width <= 0 || height <= 0) { continue; }
classIds.push_back(classIdPoint.y);
confidences.push_back(max_class_socre);
boxes.push_back(Rect(left, top, width, height));
}
}
//执行非最大抑制以消除具有较低置信度的冗余重叠框(NMS)
std::vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, CONF_THRESHOLD, NMS_THRESHOLD, nms_result);
std::vector<cv::Mat> temp_mask_proposals;
std::vector<OutputSeg> output;
Rect holeImgRect(0, 0, src.cols, src.rows);
for (int i = 0; i < nms_result.size(); ++i) {
int idx = nms_result[i];
OutputSeg result;
result.id = classIds[idx];
result.confidence = confidences[idx];
result.box = boxes[idx]& holeImgRect;
output.push_back(result);
//temp_mask_proposals.push_back(picked_proposals[idx]);
}
end = std::chrono::system_clock::now();
std::cout << "后处理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count()<< "ms" << std::endl;
DrawPred(src, output);
cv::imshow("output1.jpg", src);
auto end0 = std::chrono::system_clock::now();
以上代码为核心部分,如要学习请参考完整代码。
三、实时目标检测展示
目前在jetson tx2上运行fps为20左右,后续会继续优化。主要是图像的预处理部分暂时没有使用cuda加速,使用cuda加速之后fps应该能达到30。