【模型部署】c++调用tensorRT的模型（engine）

将分割模型就行腾搜人RT转化后得到engine，该博客主要是针对c++调用tensorRT的模型文件engine

1. 框架

文件主要分为两个文件：一个是main_tensorRT（exe），用于调用部署好的AI模型（dll）
一个是segmentationModel（dll），将tensorRT推理的程序封装成dll，用来进行调用。

2. main_tensorRT（exe）

2.1 LoadCathodeHeadEngine（读取模型）

2.2 getCathodeBiImg（运行模型）

2.3 输出结果

// Xray_test.cpp : 定义控制台应用程序的入口点。
#define _AFXDLL
#include <iomanip>
#include <string> 
#include <fstream> 
#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include <io.h>
#include "segmentationModel.h"


#define GrayScale16 65536
#define WIDTH 512
#define HEIGHT 512
#define CLWIDTH 640
#define CLHEIGHT 384


MODELDLL predictClasse;
#pragma comment(lib, "..//x64//Release//segmentationModel.lib")

using namespace cv;
using namespace std;


void getCathodeBiImg(const Mat& ROIImg, Mat& cathodeImg, int& gap);

int main()
{
    
	string filenameStrHead = "unetCathodeHead.engine";
	predictClasse.LoadCathodeHeadEngine("D:\\XrayParameters\\saveEngine\\" + filenameStrHead);
	string filenameStrTail = "unetCathodeTail.engine";
	predictClasse.LoadCathodeTailEngine("D:\\XrayParameters\\saveEngine\\" + filenameStrTail);
	string filenameStrAnode = "unetCathodeLine.engine";
	predictClasse.LoadCathodeLineEngine("D:\\XrayParameters\\saveEngine\\" + filenameStrAnode);

	string path= "D:\\Users\\6536\\Desktop\\AI_Detect\\mask\\1.png";
	Mat pre = imread(path, IMREAD_ANYDEPTH);

	int cathodeGap = 1;//分块数量
	Mat cathodeMask = cv::Mat::zeros(pre.size(), CV_8UC1);
	getCathodeBiImg(pre, cathodeMask, cathodeGap);//预测模块
	cv::imshow("img", cathodeMask);
	cv::waitKey(0);

	return 0;
}




void getCathodeBiImg(const Mat& ROIImg, Mat& cathodeImg, int& gap)
{
    

	//切割
	cv::Rect rect, rect1;
	for (int k = 0; k < gap; k++)
	{
    
		rect.x = k * ROIImg.cols / gap;
		rect.y = 0;
		rect.width = ROIImg.cols / gap;
		rect.height = ROIImg.rows;

		rect1.x = k * ROIImg.cols / gap;
		rect1.y = 0;
		rect1.width = ROIImg.cols / gap;
		rect1.height = ROIImg.rows;


		rect.x = max(0, int(rect.x));
		rect1.x = max(0, int(rect.x));


		if (ROIImg.cols % 2 != 0 && k == 0)
		{
    
			rect.width = rect.width + 61;
			rect1.width = rect1.width + 61;
		}
		else if (k != gap - 1)
		{
    
			rect.width = rect.width + 60;
			rect1.width = rect1.width + 60;
		}
		if (rect.x + rect.width > ROIImg.cols)
		{
    
			rect.width = ROIImg.cols - rect.x;
			rect1.width = ROIImg.cols - rect1.x;
		}

		
		Mat cathodeLocImg;
		predictClasse.batteryCathodeHeadPredict(ROIImg(rect), cathodeLocImg, 1, 1, WIDTH, HEIGHT);
		bitwise_or(cathodeImg(rect1), cathodeLocImg, cathodeImg(rect1));
	}
	threshold(cathodeImg, cathodeImg, 0, 255, THRESH_BINARY);
}

3. segmentationModel（dll）

3.1 LoadCathodeHeadEngine（加载模型函数）

3.2 batteryCathodeHeadPredict（模型推理）

#include "pch.h"
#include "segmentationModel.h"
#include "cuda_runtime_api.h"
#include "logging.h"
#include "common.hpp"
#include "calibrator.h"
#define USE_FP16 // comment out this if want to use FP32
#define DEVICE 0 // GPU id
#define CONF_THRESH 0.45 // 0.5
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;
//IRuntime* runtimeROI;
//ICudaEngine* engineROI;
//IExecutionContext* contextROI;
IRuntime* runtimeCathodeHead;
ICudaEngine* engineCathodeHead;
IExecutionContext* contextCathodeHead;

IRuntime* runtimeCathodeTail;
ICudaEngine* engineCathodeTail;
IExecutionContext* contextCathodeTail;

IRuntime* runtimeCathodeLine;
ICudaEngine* engineCathodeLine;
IExecutionContext* contextCathodeLine;


MODELDLL::MODELDLL()
{
    

}
MODELDLL::~MODELDLL()
{
    
   
}


bool MODELDLL::LoadCathodeHeadEngine(const std::string& engineName)
{
    
    //step1：定义打开文件的方式
    std::ifstream file(engineName, std::ios_base::out | std::ios_base::binary);// 如果想以输入方式(只写)打开，就用ifstream来定义;// 如果想以输出方式(只读)打开，就用ofstream来定义;// 如果想以输入/输出方式来打开，就用fstream来定义。
    if (!file.good())
    {
    
        return false;
    }

    //step2：定义 trtmodelstream
    char* trtModelStream = nullptr;

    //step3：获取文件大小
    size_t size = 0;
    file.seekg(0, file.end);
    size = file.tellg();

    //step4：回到文件的开头
    file.seekg(0, file.beg);

    //step5：trtModelStream被定义为一个[“引擎文件大小”]大小的字符 
    trtModelStream = new char[size];
    assert(trtModelStream);

    //step6：rtrModelStream读取文件的内容
    file.read(trtModelStream, size);

    //step7：关闭文件
    file.close();

    //step8：获取模型文件内容
    std::vector<std::pair<char*, size_t>> FileContent;
    FileContent.push_back(std::make_pair(nullptr, 0));
    FileContent[0] = std::make_pair(trtModelStream, size);
    assert(FileContent != nullptr);
   
    //step9:创建runtime
    runtimeCathodeHead = createInferRuntime(gLogger);
    assert(runtime != nullptr);

    //step10:反序列化创建engine
    engineCathodeHead = runtimeCathodeHead->deserializeCudaEngine(FileContent[0].first, FileContent[0].second, nullptr);
    assert(engine != nullptr);

    //step11:创建context,创建一些空间来存储中间激活值
    contextCathodeHead = engineCathodeHead->createExecutionContext();
    assert(context != nullptr);
	return true;
}


bool MODELDLL::batteryCathodeHeadPredict(const Mat& src, Mat& dst,  const int& channel, const int& classe, const int& width, const int& height)
{
    
    cudaSetDevice(DEVICE);
    if (src.empty())
        return false;
    //参数初始化
    float *data = new float[channel * width * height];
    float* prob = new float[classe * width * height];
    cv::Mat primg = src.clone();
    cv::resize(primg, primg, cv::Size( width,  height));
    normalizeImg(primg, data);

    //step1:使用这些索引，创建buffers指向 GPU 上输入和输出缓冲区
    const ICudaEngine& engine = (* contextCathodeHead).getEngine();
    assert(engine.getNbBindings() == 2);

    //step2：分配buffers空间,为输入输出开辟GPU显存。Allocate GPU memory for Input / Output data
    void* buffers[2];
    const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
    const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
    CHECK(cudaMalloc(&buffers[inputIndex], channel * width * height*sizeof(float)));
    CHECK(cudaMalloc(&buffers[outputIndex], classe * width * height * sizeof(float)));
 
    step3：创建cuda流
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));
  
    step4：从CPU到GPU----拷贝input数据
    CHECK(cudaMemcpyAsync(buffers[inputIndex], data, channel * width * height * sizeof(float), cudaMemcpyHostToDevice, stream));
   
    step5: 启动cuda核计算
    (*contextCathodeHead).enqueue(1, buffers, stream, nullptr);
  
    step6：从GPU到CPU----拷贝output数据
    CHECK(cudaMemcpyAsync(prob, buffers[outputIndex], classe *width * height * sizeof(float), cudaMemcpyDeviceToHost, stream));
    cudaStreamSynchronize(stream);
   
    step7：获取Mat格式的输出结果，并调整尺寸
    float* mask = prob;
    dst = cv::Mat(height, width, CV_8UC1);
    uchar* ptmp = NULL;
    for (int i = 0; i < height; i++)
    {
    
        ptmp = dst.ptr<uchar>(i);
        for (int j = 0; j < width; j++)
        {
    
            float* pixcel = mask + i * width + j;
            if (*pixcel > CONF_THRESH)
            {
    
                ptmp[j] = 255;
            }
            else
            {
    
                ptmp[j] = 0;
            }
        }
    }
    cv::resize(dst, dst, src.size(), INTER_NEAREST);
   
    step8：释放资源
    cudaStreamDestroy(stream);
    CHECK(cudaFree(buffers[inputIndex]));
    CHECK(cudaFree(buffers[outputIndex]));
    delete prob;
    delete data;
}

4. 输出结果

5. 完整项目

https://download.csdn.net/download/qq_44747572/88027027