部署：针对yolov5网络输出的结果，对数据做后处理以输出目标检测结果

针对yolov5网络输出的结果，对数据做后处理以输出目标检测结果

engine网络后处理

包含：

​ （1）传入一张图片转为需要的格式（参看：(63条消息) 对图像做前处理，以便按照yolov5的engine输入格式输入网络_爱吃油淋鸡的莫何的博客-CSDN博客）

​ （2）调用engine进行推理了

​ （3）对输出的后处理

​ （4）输出结果绘图

需要注意的是：

1 pytorch的pt文件转.onnx文件的时候涉及batchsize值，onnx2engine的时候也需要设置batchsize值，infer推理的时候也有batchsize参数。上述三处的batchsize的值需要一致，否则会出现【Cuda failure: 700 已放弃 (核心已转储)】的错误提示。

2 代码仅演示了一张固定图片的推理操作。

代码：

1 CMakeLists.txt

# CMakeLists.txt
cmake_minimum_required(VERSION 2.6)

project(yolo)

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}/postprocess.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)

2 postprocess.cpp

 // yolo_onnx2engine.cpp
 // ref1 : https://blog.csdn.net/weixin_43863869/article/details/124614334
 // ref2 : https://www.cnblogs.com/tangjunjun/p/16639361.html
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#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>


// onnx转换头文件
#include "NvOnnxParser.h"
using namespace nvonnxparser;
using namespace nvinfer1;
using namespace std;


#define CHECK(status) \
    do\
    {\
        auto ret = (status);\
        if (ret != 0)\
        {\
            std::cerr << "Cuda failure: " << ret << std::endl;\
            abort();\
        }\
    } while (0)


// stuff we know about the network and the input/output blobs
static const int INPUT_H = 960;
static const int INPUT_W = 960;
static const int cls_num = 15;
static const unsigned int batchSize = 1;
static const int anchor_output_num = 56700;   // 不同输入尺寸anchor:640-->25200 | 960-->56700
static const int OUTPUT_SIZE = batchSize * anchor_output_num *(cls_num+5); //1000 * sizeof(Detection) / sizeof(float) + 1;

const char* INPUT_BLOB_NAME = "images";
const char* OUTPUT_BLOB_NAME = "output";
const char* TrtFileName = "/home/onnx2engine/yolov5_5.0/model_onnx/yolov5s.engine";

//构建Logger
class Logger : public ILogger
{
    void log(Severity severity, const char* msg) noexcept override
    {
        // suppress info-level messages
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} gLogger;

//********************************************** 使用一张图片查看推理结果 **********************************************//

struct  Detection {
    //center_x center_y  w h
    float bbox[4];
    float conf;  // bbox_conf * cls_conf
    int class_id;
    int index;
};

struct Bbox {
    int x;
    int y;
    int w;
    int h;
};

//加工图片变成拥有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);
        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++;
            }
        }
    }
}

float iou(Bbox box1, Bbox box2) {
    /*  
    iou=交并比
    */
    int x1 = max(box1.x, box2.x);
    int y1 = max(box1.y, box2.y);
    int x2 = min(box1.x + box1.w, box2.x + box2.w);
    int y2 = min(box1.y + box1.h, box2.y + box2.h);
    int w = max(0, x2 - x1);
    int h = max(0, y2 - y1);
    float over_area = w * h;
    return over_area / (box1.w * box1.h + box2.w * box2.h - over_area);
}

int get_max_index(vector<Detection> pre_detection) {
    //返回最大置信度值对应的索引值
    int index;
    float conf;
    if (pre_detection.size() > 0) {
        index = 0;
        conf = pre_detection.at(0).conf;
        for (int i = 0; i < pre_detection.size(); i++) {
            if (conf < pre_detection.at(i).conf) {
                index = i;
                conf = pre_detection.at(i).conf;
            }
        }
        return index;
    }
    else {
        return -1;
    }
}

bool judge_in_lst(int index, vector<int> index_lst) {
    //若index在列表index_lst中则返回true，否则返回false
    if (index_lst.size() > 0) {
        for (int i = 0; i < index_lst.size(); i++) {
            if (index == index_lst.at(i)) {
                return true;
            }
        }
    }
    return false;
}

vector<int> nms(vector<Detection> pre_detection, float iou_thr)
{
    /*
    返回需保存box的pre_detection对应位置索引值
    */
    int index;
    vector<Detection> pre_detection_new;
    //Detection det_best;
    Bbox box_best, box;
    float iou_value;
    vector<int> keep_index;
    vector<int> del_index;
    bool keep_bool;
    bool del_bool;

    if (pre_detection.size() > 0) {
        pre_detection_new.clear();
        // 循环将预测结果建立索引
        for (int i = 0; i < pre_detection.size(); i++) {
            pre_detection.at(i).index = i;
            pre_detection_new.push_back(pre_detection.at(i));
        }
        //循环遍历获得保留box位置索引-相对输入pre_detection位置
        while (pre_detection_new.size() > 0) {
            index = get_max_index(pre_detection_new);
            if (index >= 0) {
                keep_index.push_back(pre_detection_new.at(index).index); //保留索引位置

                // 更新最佳保留box
                box_best.x = pre_detection_new.at(index).bbox[0];
                box_best.y = pre_detection_new.at(index).bbox[1];
                box_best.w = pre_detection_new.at(index).bbox[2];
                box_best.h = pre_detection_new.at(index).bbox[3];

                for (int j = 0; j < pre_detection.size(); j++) {
                    keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
                    del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
                    if ((!keep_bool) && (!del_bool)) { //不在keep_index与del_index才计算iou
                        box.x = pre_detection.at(j).bbox[0];
                        box.y = pre_detection.at(j).bbox[1];
                        box.w = pre_detection.at(j).bbox[2];
                        box.h = pre_detection.at(j).bbox[3];
                        iou_value = iou(box_best, box);
                        if (iou_value > iou_thr) {
                            del_index.push_back(j); //记录大于阈值将删除对应的位置
                        }
                    }

                }
                //更新pre_detection_new
                pre_detection_new.clear();
                for (int j = 0; j < pre_detection.size(); j++) {
                    keep_bool = judge_in_lst(pre_detection.at(j).index, keep_index);
                    del_bool = judge_in_lst(pre_detection.at(j).index, del_index);
                    if ((!keep_bool) && (!del_bool)) {
                        pre_detection_new.push_back(pre_detection.at(j));
                    }
                }
            }
        }
    }

    del_index.clear();
    del_index.shrink_to_fit();
    pre_detection_new.clear();
    pre_detection_new.shrink_to_fit();

    return  keep_index;

}

void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
    const ICudaEngine& engine = context.getEngine();
    // 指向要传递给引擎的输入和输出设备缓冲区的指针.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(engine.getNbBindings() == 2);
    void* buffers[2];
    // 为了绑定缓冲区，我们需要知道输入和输出张量的名称.注意，确保索引小于IEngine:：getNbBindings（
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);  // "images"
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);// "output"
    //const int inputIndex = 0;
    //const int outputIndex = 1;
    // Create GPU buffers on device    
    cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float));
    cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float));
    // Create stream
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));
    // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
    // DMA 将批处理数据输入到设备，异步推断批处理，并将DMA输出回主机
    CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
    context.enqueue(batchSize, buffers, stream, nullptr);
    cout<<"1"<<endl;
    CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
    cout<<"2"<<endl;
    cudaStreamSynchronize(stream);
    cout<<"3"<<endl;
    // Release stream and buffers
    cudaStreamDestroy(stream);
    CHECK(cudaFree(buffers[inputIndex]));
    CHECK(cudaFree(buffers[outputIndex]));
}

vector<Detection> postprocess(float* prob, float conf_thr = 0.2, float nms_thr = 0.4) {
    /*
    #####################此函数处理一张图预测结果#########################
    prob为[x y w h  score  multi-pre] 如80类-->(1,anchor_num,85)
    (1,anchor_num,85)理解为 如果batchsize=1，则输出数据三维，
    第二维和第三维理解为每行是[x y w h  bbox_conf  cls0_conf cls1_conf cls2_conf ... ]，共anchor_num行。
    实际上 prob 是输出结果第一个数据的指针(最终的所有数据像一个向量一样依次放在指针指向的及其后面的地址中)
    */
    vector<Detection> pre_results;
    vector<int> nms_keep_index;
    vector<Detection> results;
    bool keep_bool;
    Detection pre_res;
    float conf;
    int tmp_idx;
    float tmp_cls_score;
    for (int i = 0; i < anchor_output_num; i++) {
        tmp_idx = i * (cls_num + 5);
        pre_res.bbox[0] = prob[tmp_idx + 0];  // 从第一个指针开始，每（cls_num+5）个数据中，前4位对应x,y,w,h
        pre_res.bbox[1] = prob[tmp_idx + 1];
        pre_res.bbox[2] = prob[tmp_idx + 2];
        pre_res.bbox[3] = prob[tmp_idx + 3];
        conf = prob[tmp_idx + 4];  // 是为目标的置信度
        tmp_cls_score = prob[tmp_idx + 5] * conf; 
        pre_res.class_id = 0;
        pre_res.conf = 0;
        // 这个过程相当于从除了前面5列，在后面的cla_num个数据中找出score最大的值作为pre_res.conf，对应的列作为类id
        for (int j = 1; j < cls_num; j++) {     
            tmp_idx = i * (cls_num + 5) + 5 + j; //获得对应类别索引
            if (tmp_cls_score < prob[tmp_idx] * conf)
            {
                tmp_cls_score = prob[tmp_idx] * conf;
                pre_res.class_id = j;
                pre_res.conf = tmp_cls_score;
            }
        }
        if (conf >= conf_thr) {
            pre_results.push_back(pre_res);
        }
    }

    //使用nms，返回对应结果的索引
    nms_keep_index=nms(pre_results,nms_thr);

    // 茛据nms找到的索引，将结果取出来作为最终结果
    for (int i = 0; i < pre_results.size(); i++) {
        keep_bool = judge_in_lst(i, nms_keep_index);
        if (keep_bool) {
            results.push_back(pre_results.at(i));
        }
    }

    pre_results.clear();
    pre_results.shrink_to_fit();
    nms_keep_index.clear();
    nms_keep_index.shrink_to_fit();

    return results;
}

cv::Mat draw_rect(cv::Mat image, vector<Detection> results) 
{   
    float x;
    float y;
    float y_tmp;
    float w;
    float h;
    string info;

    cv::Rect rect;
    for (int i = 0; i < results.size(); i++) 
    {
        
        x = results.at(i).bbox[0];
        y= results.at(i).bbox[1];
        w= results.at(i).bbox[2];
        h=results.at(i).bbox[3];
        x = (int)(x - w / 2);
        y = (int)(y - h / 2);
        w = (int)w;
        h = (int)h;
        info = "id:";
        info.append(to_string(results.at(i).class_id));
        info.append(" s:");
        info.append(  to_string((int)(results.at(i).conf*100)  )   );
        info.append("%");
        rect= cv::Rect(x, y, w, h);
        
        cv::rectangle(image, rect, cv::Scalar(0, 255, 0), 1, 1, 0);
        cv::putText(image, info, cv::Point(x, y), cv::FONT_HERSHEY_SIMPLEX, 0.4, cv::Scalar(0, 255, 0), 0.4, 1, false);    
    }
    return image;
}

// infer 
int main(int argc, char** argv)
{ 
    // cout<<"argv = "<<argv<<endl;
    // cout<<"argv[1] = "<<argv[1]<<endl;
    //加载engine引擎
    char* trtModelStream{ nullptr };
    size_t size{ 0 };
    std::ifstream file(TrtFileName, std::ios::binary); // 以二进制模式打开文件
    if (file.good()) {
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream = new char[size];
        assert(trtModelStream);
        file.read(trtModelStream, size);
        file.close();
    }

    //反序列为engine，创建context
    IRuntime* runtime = createInferRuntime(gLogger);
    assert(runtime != nullptr);
    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
    assert(engine != nullptr);
    IExecutionContext* context = engine->createExecutionContext();
    assert(context != nullptr);
    delete[] trtModelStream;

    //*********************推理-循环推理*********************//

    float time_read_img = 0.0;
    float time_infer = 0.0;
    float prob[OUTPUT_SIZE];
    vector<Detection> results;

    for (int i = 0; i < 1000; i++) {
        
        // 处理图片为固定输出
        auto start = std::chrono::system_clock::now();  //时间函数
        std::string path = "/home/sources/data220729_0467.jpg";
        std::cout << "img_path=" << path << endl;
        static float data[3 * INPUT_H * INPUT_W];
        cv::Mat img = cv::imread(path);
        
        ProcessImage(img, data);
        auto end = std::chrono::system_clock::now();
        time_read_img = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() + time_read_img;
        cout<<"......ProcessImage结束 。。。"<<endl;
        
        //Run inference
        start = std::chrono::system_clock::now();  //时间函数
        doInference(*context, data, prob, batchSize);
        end = std::chrono::system_clock::now();
        time_infer = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() + time_infer;
        std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
        cout<<"......doInference结束 。。。"<<endl;

        //输出后处理
        std::cout <<"prob="<<prob << std::endl;
        results.clear();
        results=postprocess(prob, 0.3,  0.4);

        cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), 0, 0, cv::INTER_LINEAR);
        img=draw_rect(img, results);

        //cv::imshow("www", img);
        //cv::waitKey(0);
        cv::imwrite("./ftest.jpg", img);

        cout << "ok" << endl;
    }

    std::cout << "C++ 2engine" << "mean read img time =" << time_read_img / 1000 << "ms\t" << "mean infer img time =" << time_infer / 1000 << "ms" << std::endl;

    // Destroy the engine
    context->destroy();
    engine->destroy();
    runtime->destroy();

    return 0;

}

若有描述或理解错误，请指正，非常感谢！

相关代码链接：链接：https://pan.baidu.com/s/1IUPyVb15PpWpnJ4lrLVFNQ