TensorST 部署YOLOv5 C++处理

蓝天巨人

已于 2023-06-07 11:11:13 修改

阅读量675

点赞数 2

分类专栏： TensorRT学习文章标签： YOLO c++

于 2023-03-02 16:48:30 首次发布

本文链接：https://blog.csdn.net/qq_44632658/article/details/129249906

版权

TensorRT学习专栏收录该内容

6 篇文章 0 订阅

订阅专栏

简单介绍

本文使用的TensorST进行模型的推理工作，为了说明整个具体的处理过程的，本文针对具体的处理过程进行记录说明

构建的生成器

为了进行模型的上下文推理，若要创建对应生成器，必须首先的进行的实例化ILogger接口，利用的ILogger进行输出模型处理过程的日志。

class Logger:public ILogger{
    public:
        void log(Severity severity, const char* msg) noexcept override{
            if(severity<Serverity::kWARNING){
                std::cerr << msg << std::endl;
            }
        }
}logger;

Ibuilder  builder = createInferBuilder(logger);

new-tensorrtx: https://github.com/wang-xinyu/tensorrtx

yolov5s.trt与yolov5s.engine是一样的，只是后缀名不同。

本文使用YOLOv5的模型到处ONNX模型进行YOLOv5的tensorRT进行推理，根据整个数据的处理流程包含以下的几个部分。

图像的预处理
模型推理
图像结果

1，进行模型推理的前处理

/ 进行图像的预处理改变图像大小
cv::Mat YOLOv5::resize_image(cv::Mat srcimg,int *newh,int *neww,int *top,int *left){
    int src_h = srcimg.rows;
    int src_w = srcimg.cols;
    *newh = this->in_p_height;
    *neww = this->in_p_width;
    Mat dst_img;
    if(this->keep_ratio&&src_h!=src_w){
        float hw_ratio = (float)src_h / (float)src_w;
        if(hw_ratio>1){
            *newh = this->in_p_height;
            *neww = int(this->in_p_height/hw_ratio);
            resize(srcimg,dst_img,Size(*newh,*neww),INTER_AREA);
            *left = int((this->in_p_width-*neww)*0.5);
            copyMakeBorder(dst_img,dst_img,0,0,*left,this->in_p_width-*neww-*left,BORDER_CONSTANT,114);

        }else{
            *newh = (int)this->in_p_height * hw_ratio;
			*neww = this->in_p_width;
			resize(srcimg, dst_img, Size(*neww, *newh), INTER_AREA);
			*top = (int)(this->in_p_height - *newh) * 0.5;
			copyMakeBorder(dst_img, dst_img, *top, this->in_p_height - *newh - *top, 0, 0, BORDER_CONSTANT, 114);
		
        }
    }else{
        resize(srcimg, dst_img, Size(*neww, *newh), INTER_AREA);
    }
    return dst_img;
}

// 输入图像的归一化处理
void YOLOv5::normalize_(Mat &img){
    int row = img.rows;
    int col = img.cols;
    this->input_image_.resize(row*col*img.channels());
    // bgr
    for(int c = 0; c < 3;c++){
        for(int i= 0;i<row;i++){
            for(int j=0;j<col;j++){
                float pix = img.ptr<uchar>(i)[j*3+2-c];
                this->input_image_[c*row*col+i*col+j] = pix/255.0;
            }
        }
    }
}


void YOLOv5::nms(vector<BoxInfo> &input_boxes){
    sort(input_boxes.begin(), input_boxes.end(),[](BoxInfo a, BoxInfo b){ return a.score>b.score; });
    vector<float> vArea(input_boxes.size());
    for(int i = 0; i < input_boxes.size();i++){
        vArea[i] = (input_boxes[i].x2-input_boxes[i].x1+1)*(input_boxes[i].y2-input_boxes[i].y1+1);
    }
    vector<bool> isSuppressed(input_boxes.size(),false);
    for(int i = 0; i < input_boxes.size();i++){
        if (isSuppressed[i]) { continue; }
        for (int j = i + 1; i < input_boxes.size(); j++)
        {
            if (isSuppressed[j]) { continue; }
            if (isSuppressed[j]) { continue; }
			float xx1 = max(input_boxes[i].x1, input_boxes[j].x1);
			float yy1 = max(input_boxes[i].y1, input_boxes[j].y1);
			float xx2 = min(input_boxes[i].x2, input_boxes[j].x2);
			float yy2 = min(input_boxes[i].y2, input_boxes[j].y2);
            float w = max(0.0f, xx2 - xx1 + 1);
			float h = max(0.0f, yy2 - yy1 + 1);
			float inter = w * h;	// 交集
			if(input_boxes[i].label == input_boxes[j].label)
			{
				float ovr = inter / (vArea[i] + vArea[j] - inter);  // 计算iou
				if (ovr >= this->nms_threshold)
				{
					isSuppressed[j] = true;
				}
			}	
        }
        
    }
    int idx_t = 0;
       // remove_if()函数 remove_if(beg, end, op) //移除区间[beg,end)中每一个“令判断式:op(elem)获得true”的元素
	input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfo& f) { return isSuppressed[idx_t++]; }), input_boxes.end());


}

2 模型的处理过程

利用已有的模型进行对应的处理，本文中提供了针对ONNX模型的序列化处理的工作，相关的处理流出如下所示。

void YOLOv5::loadOnnx(const std::string strModelName)
{
    TRT_Logger gLogger;   // 日志
    //根据tensorrt pipeline 构建网络
    IBuilder* builder = createInferBuilder(gLogger);    // 网络元数据,用于搭建网络入口 
    builder->setMaxBatchSize(1);   // batchsize
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);  // 显式批处理
    INetworkDefinition* network = builder->createNetworkV2(explicitBatch);                      // 定义模型
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);              // 使用nvonnxparser 定义一个可用的onnx解析器
    parser->parseFromFile(strModelName.c_str(), static_cast<int>(ILogger::Severity::kWARNING));   // 解析onnx
    // 使用builder对象构建engine
    IBuilderConfig* config = builder->createBuilderConfig();   // 
    // 特别重要的属性是最大工作空间大小
    config->setMaxWorkspaceSize(1ULL << 30);                   // 分配内存空间
    m_CudaEngine = builder->buildEngineWithConfig(*network, *config);    // 来创建一个 ICudaEngine 类型的对象，在构建引擎时，TensorRT会复制权重

    std::string strTrtName = strModelName;
    size_t sep_pos = strTrtName.find_last_of(".");
    strTrtName = strTrtName.substr(0, sep_pos) + ".trt";  // 
    IHostMemory *gieModelStream = m_CudaEngine->serialize();    // 将引擎序列化
    std::string serialize_str;     // 
    std::ofstream serialize_output_stream;
    serialize_str.resize(gieModelStream->size()); 
    // memcpy内存拷贝函数 ，从源内存地址的起始位置开始拷贝若干个字节到目标内存地址中
    memcpy((void*)serialize_str.data(),gieModelStream->data(),gieModelStream->size()); 
    serialize_output_stream.open(strTrtName.c_str());  
    serialize_output_stream<<serialize_str;     // 将引擎序列化数据转储到文件中
    serialize_output_stream.close();   
    m_CudaContext = m_CudaEngine->createExecutionContext();    //执行上下文用于执行推理
	// 使用一次，销毁parser，network, builder, and config 
    parser->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
}

同样针对现有的TRT模型或者engine模型进行的反序列化推理的工作

  // 进行对应的引擎的序列化
    m_CudaRuntime=createInferRuntime(gLogger);
    ifstream fin(strName.c_str());
    std::string cached_engine = "";
    while (fin.peek()!=EOF)
    {
        stringstream buffer;
        buffer<<fin.rdbuf();
        cached_engine.append(buffer.str());
    }
    fin.close();
    // 将对应的模型输入到缓冲区后，将其进行的反序列化
    m_CudaEngine = m_CudaRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(),nullptr);
    // 针对engine进行的上下文内容的执行
    m_CudaExecutionContext = m_CudaEngine->createExecutionContext();
    m_CudaRuntime->destroy();

进行模型的检测和推理动作

void YOLOv5::detect(Mat& frame){
    int newh = 0;
    int neww = 0;
    int padh = 0;
    int padw = 0;
    Mat dstimg=this->resize_image(frame,&newh,&neww,&padh,&padw);
    this->normalize_(dstimg);
    // 定义一个对应的输入举证
    array<int64_t,4> input_shape_{1,3,this->in_p_height,this->in_p_width};
    // 进行对应通道类型的转化
    cvtColor(dstimg,  dstimg, cv::COLOR_BGR2RGB);
    cv::Mat m_Normalized;
    dstimg.convertTo(m_Normalized,CV_32FC3,1/255.);
    // 进行对应的通道分离
    cv::split(m_Normalized,m_InputWrapper);
    // 创建CUDA流，推理时TensorRT是异步的，因此将CUDA流
    cudaStreamCreate(&m_CudaStream);
    auto ret = cudaMemcpyAsync(m_ArrayDevMemory[m_iInputIndex], m_ArrayHostMemory[m_iInputIndex], 
                                                                m_ArraySize[m_iInputIndex], cudaMemcpyHostToDevice, m_CudaStream); 
	// 进行异步执行，将内核排入到CUDA流中
    auto ret1 = m_CudaExecutionContext->enqueueV2(m_ArrayDevMemory, m_CudaStream, nullptr);    // TensorRT 执行通常是异步的，因此将内核排入 CUDA 流：
	ret = cudaMemcpyAsync(m_ArrayHostMemory[m_iOutputIndex], m_ArrayDevMemory[m_iOutputIndex], m_ArraySize[m_iOutputIndex], 
                            cudaMemcpyDeviceToHost, m_CudaStream); //输出传回给CPU，数据从显存到内存
    ret = cudaStreamSynchronize(m_CudaStream);
    float * pdata = (float *)m_ArrayHostMemory[m_iOutputIndex];
    std::vector<BoxInfo> generate_boxes;
    float ratioh = (float)frame.rows / newh, ratiow = (float)frame.cols / neww;
    for(int i = 0; i < m_iBoxNums;++i){
        int index = i * (m_iClassNums + 5);      // prob[b*num_pred_boxes*(classes+5)]  
		float obj_conf = pdata[index + 4];  // 置信度分数
        if (obj_conf > this->obj_threshold)  // 大于阈值
		{
			float* max_class_pos = std::max_element(pdata + index + 5, pdata + index + 5 + m_iClassNums);   //
			(*max_class_pos) *= obj_conf;   // 最大的类别分数*置信度
			if ((*max_class_pos) > this->conf_threshold) // 再次筛选
			{ 
				//const int class_idx = classIdPoint.x;
				float cx = pdata[index];  //x
				float cy = pdata[index+1];  //y
				float w = pdata[index+2];  //w
				float h = pdata[index+3];  //h

				float xmin = (cx - padw - 0.5 * w)*ratiow;
				float ymin = (cy - padh - 0.5 * h)*ratioh;
				float xmax = (cx - padw + 0.5 * w)*ratiow;
				float ymax = (cy - padh + 0.5 * h)*ratioh;

				generate_boxes.push_back(BoxInfo{ xmin, ymin, xmax, ymax, (*max_class_pos), int(max_class_pos-(pdata + index + 5)) });
			}
		}
    }
    nms(generate_boxes);
    for (int i = 0; i <generate_boxes.size(); i++){

		int xmin = int(generate_boxes[i].x1);
		int ymin = int(generate_boxes[i].y1);
		rectangle(frame, Point(xmin, ymin), Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), Scalar(0, 0, 255), 2);
		std::string label = format("%.2f", generate_boxes[i].score);
		label = this->classes[generate_boxes[i].label] + ":" + label;
		putText(frame, label, Point(xmin, ymin - 5), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0), 1);
    }