DepthAnything(3): 基于TensorRT在NVIDIA边缘平台部署DepthAnything(pt转onnx,经trtexec转engine,tensorrt反序列化推理)

DepthAnything(1): 先跑一跑Depth Anything_depth anything离线怎么跑-CSDN博客

DepthAnything(2): 基于ONNXRuntime在ARM(aarch64)平台部署DepthAnything-CSDN博客


目录

1. 写在前面

2. 转ONNX模型

3. 使用trtexec转换

4. 基本的API介绍

5. 推理例程


1. 写在前面

        DepthAnything是一种能在任何情况下处理任何图像的简单却又强大的深度估计模型。

      TensorRT是NVIDIA推出的运行在NV平台上的推理加速框架,我们常用的TX2、Jetson Nano、Jetson Xavier NX以及Jetson Orin系列均可通过TensorRT实现推理加速。TensorRT 基于 C++ 语言开发,能够实现高性能推理,TensorRT 推理应用程序与英伟达的 GPU 结合可以提供最大的吞吐量和较低的延迟速度

2. PyTorch的pt权重转onnx

        以DepthAnything为例,使用export_onnx.py导出ONNX模型。

        转换模型的时候,需要注意设定输入输出节点的命名。

        具体转换代码,可参考上一节【DepthAnything(2): 基于ONNXRuntime在ARM(aarch64)平台部署DepthAnything-CSDN博客】。

3. 使用trtexec生成引擎(基于onnx模型)

        转换时,尽量使用sudo trtexec,避免因权限问题导致的转换失败。如下为一个简单的例子。

sudo trtexec --onnx=depth_anything_vitl14.onnx --saveEngine=depth_anything_vitl14.engine --explicitBatch --verbose

        注意,trtexec可能不在常规搜索路径里,这个时候需要使用完整的路径,如/usr/src/tensorrt/bin/trtexec。

4. 基本的API介绍

        cudaStreamCreate:创建处理流;

        createInferRuntime:创建运行时;

        deserializeCudaEngine: 反序列化推理引擎;

        createExecutionContext:创建上下文;

        getNbBindings:获取输入输出总数;

        getBindingName:获取绑定输入/输出Tensor名;

        getBindingIndex:获取输入输出对应的索引;

        bindingIsInput:判断是输入或输出Tensor;

        enqueueV2:执行前向推理(异步);

        executeV2:执行前向推理(同步);

5. 推理例程

        如下是一个使用ONNX转换得到推理引擎,配合TensorRT和OpenCV完成推理的代码。

        完整的代码工程可访问以下链接下载,下载完成后进入工程一级目录,执行sudo make即可。

        链接:https://pan.baidu.com/s/17gTSbUhn6sAr92JvWLKZcg?pwd=1234
        提取码:1234

        推理输出如下图所示。

#include "NvInfer.h"

#include "NvInferRuntime.h"

#include "cuda_runtime_api.h"

#include "logging.h"

#include "trt_utils.h"

#include <opencv2/opencv.hpp>

#include <iostream>  

#include <fstream>  

#include <vector>  

 

using namespace nvinfer1;  

Logger gLogger;


std::string engine_path;

size_t EngineSize = 0;

char *trtModelStream = NULL;

IRuntime* runtime;

ICudaEngine* engine;

IExecutionContext* context;

int batchsize = 1;

Dims input_dims;

std::string input_image;

int input_width = 518;

int input_height = 518;

float* data;

// float* DepthMap[1];

std::vector<float> DepthMap(input_width*input_height);

void* buffers[2];

std::string input_name;

std::string output_name;

int input_index = -1;

int output_index = -1;

cudaStream_t stream;

float CalculateCentrifugalDistance(cv::Mat ori_img, cv::Point2f center);

int main(int argc, char* argv[])

{

    engine_path = std::string(argv[1]);

    input_image = std::string(argv[2]);

    input_dims.nbDims = 4;

    input_dims.d[0] = batchsize;

    input_dims.d[1] = 3;

    input_dims.d[2] = input_height;

    input_dims.d[3] = input_width;

    CUDA_CHECK(cudaStreamCreate(&stream));

    /** 创建TensorRT模型的运行时. */

    runtime = createInferRuntime(gLogger);  

    if (!runtime) {  

        return 1;  

    }

    /** 读取本地引擎文件. */

    std::ifstream file(engine_path, std::ios::binary);

    if (file.good()) {

        file.seekg(0, file.end);

        EngineSize = file.tellg();

        file.seekg(0, file.beg);

        trtModelStream = new char[EngineSize];

        assert(trtModelStream);

        file.read(trtModelStream, EngineSize);

        file.close();

    }

 

    /** 反序列化、加载TensorRT引擎. */

    engine = runtime->deserializeCudaEngine(trtModelStream, EngineSize, nullptr);  

    if (!engine) {  

        return 1;  

    }  

 

    /** 创建执行上下文. */

    context = engine->createExecutionContext();  

    if (!context) {  

        return 1;  

    }  

    /** 构造输入数据缓冲. */

    data = new float[batchsize * (input_width * input_height * 3) * sizeof(float)]();

 

    /** 设置输入和输出缓冲区. */ 

    for (int i = 0; i < engine->getNbBindings(); i++) {

        auto bindingName = engine->getBindingName(i);

        auto bindingIdx = engine->getBindingIndex(bindingName);

        auto isInput = engine->bindingIsInput(bindingIdx);

        printf("Binding name: %s, binding index: %d, is input: %d\n", bindingName, bindingIdx, isInput);

        if(isInput){ ///< 输入

            CUDA_CHECK(cudaMalloc(&buffers[bindingIdx], batchsize * (input_width * input_height * 3) * sizeof(float)));

            input_index = bindingIdx;

            input_name = std::string(bindingName);

        }else{ ///< 输出

            CUDA_CHECK(cudaMalloc(&buffers[bindingIdx], batchsize * (input_width * input_height) * sizeof(float)));

            output_index = bindingIdx;

            output_name = std::string(bindingName);

        }

    }  

    cv::Mat image_mat = cv::imread(input_image);

    auto ori_h = image_mat.cols;

    auto ori_w = image_mat.rows;

    printf("Original, w: %d, h: %d\n", ori_w, ori_h);

    cv::cvtColor(image_mat, image_mat, cv::COLOR_BGR2RGB);

    cv::resize(image_mat, image_mat, {input_width, input_height}, 0.0, 0.0, cv::INTER_CUBIC);

    /** 转换RGB Planar, 归一化. */

    for (int i = 0; i < input_height * input_width; i++) { ///< HWC2CHW, 归一化

        data[i] = image_mat.at<cv::Vec3b>(i)[2] / 255.0;

        data[i + input_height * input_width] = image_mat.at<cv::Vec3b>(i)[1] / 255.0;

        data[i + 2 * input_height * input_width] = image_mat.at<cv::Vec3b>(i)[0] / 255.0;

    }

    std::chrono::time_point<std::chrono::system_clock> tic;

    std::chrono::time_point<std::chrono::system_clock> toc;

    uint64_t elapsed_time = 0;

    // while(1){

        tic = std::chrono::system_clock::now();

        CUDA_CHECK(cudaMemcpyAsync(buffers[input_index], data, batchsize * (input_width * input_height * 3) * sizeof(float), cudaMemcpyHostToDevice, stream));

        context->enqueueV2(buffers, stream, nullptr);

        // context->executeV2(buffers, stream, nullptr);

        context->setBindingDimensions(input_index, input_dims); ///< [SAI-KEY] Set input dimensions.

        CUDA_CHECK(cudaMemcpyAsync(DepthMap.data(),  buffers[output_index],  batchsize * (input_width * input_height) * sizeof(float), cudaMemcpyDeviceToHost, stream));

        toc = std::chrono::system_clock::now();

        elapsed_time = std::chrono::duration_cast<std::chrono::milliseconds>(toc - tic).count();

        printf("Infer frame takes: %d ms\n", elapsed_time);

    // }

    cv::Mat output_tensor(DepthMap);

    output_tensor =output_tensor.reshape(1, {input_width, input_height});

    double minVal, maxVal;

    cv::minMaxLoc(output_tensor, &minVal, &maxVal);

    printf("minVal: %f,  maxVal: %f\n", minVal, maxVal);

    output_tensor.convertTo(output_tensor, CV_32F);

    if (minVal != maxVal) {

        output_tensor = (output_tensor - minVal) / (maxVal - minVal);

       

    }

    output_tensor *= 255.0;

    output_tensor.convertTo(output_tensor, CV_8UC1);

    // cv::applyColorMap(output_tensor, output_tensor, cv::COLORMAP_HOT);

    // cv::resize(output_tensor, output_tensor, {input_height, input_width}, 0.0, 0.0, cv::INTER_CUBIC);

    /** 计算均值,方差. */

    cv::Mat mean_mat, stddev_mat;

    cv::meanStdDev(output_tensor, mean_mat, stddev_mat, cv::noArray());

    printf("mean_mat.shape(%d, %d, %d), mean_mat.type(%d)\n", mean_mat.channels(), mean_mat.rows, mean_mat.cols, mean_mat.type());

    printf("mean_mat: %f\n", mean_mat.at<double>(0, 0));

    printf("stddev_mat: %f\n", stddev_mat.at<double>(0, 0));

    uint8_t mean_int = static_cast<uint8_t>(mean_mat.at<double>(0, 0));

    uint8_t stdd_int = static_cast<uint8_t>(stddev_mat.at<double>(0, 0));

    printf("mean_mat: %d\n", mean_int);

    printf("stddev_mat: %d\n", stdd_int);

    cv::Mat depth_map = output_tensor;

    cv::applyColorMap(depth_map, depth_map, cv::COLORMAP_HOT);

#if 1

    /**

     * 二值化

     * 阈值设置为均值的1.7倍, 1.7倍是一个经验值。

     */

    cv::threshold(output_tensor, output_tensor, mean_int*1.8, 255, cv::THRESH_BINARY);

    std::vector<std::vector<cv::Point>> first_filter_contours;

    std::vector<cv::Vec4i> first_filter_hierarcy;

    std::vector<cv::RotatedRect> first_filter_box; ///< 最小旋转外接矩形集合

    first_filter_contours.clear();

    first_filter_hierarcy.clear();

    /** 查找连通域. */

    cv::findContours(output_tensor, first_filter_contours, first_filter_hierarcy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_NONE);

    /** 存储每一个连通域的最小外接矩形四个顶点. */

    first_filter_box.resize(first_filter_contours.size());

   

    /** 分别查找连通域的顶点. */

    cv::Point2f rect[4];

    for(int i=0; i<first_filter_contours.size(); i++){

        first_filter_box[i] = cv::minAreaRect(cv::Mat(first_filter_contours[i]));  ///< 最小旋转外接矩形

        first_filter_box[i].points(rect);  ///< 把最小外接矩形四个端点复制给rect数组

        cv::line(depth_map, cv::Point((int)rect[0].x, (int)rect[0].y), cv::Point((int)rect[1].x, (int)rect[1].y), cv::Scalar(29, 128, 255), 2);

        cv::line(depth_map, cv::Point((int)rect[1].x, (int)rect[1].y), cv::Point((int)rect[2].x, (int)rect[2].y), cv::Scalar(29, 128, 255), 2);

        cv::line(depth_map, cv::Point((int)rect[2].x, (int)rect[2].y), cv::Point((int)rect[3].x, (int)rect[3].y), cv::Scalar(29, 128, 255), 2);

        cv::line(depth_map, cv::Point((int)rect[3].x, (int)rect[3].y), cv::Point((int)rect[0].x, (int)rect[0].y), cv::Scalar(29, 128, 255), 2);

    }

#endif

   

    cv::resize(depth_map, depth_map, {ori_h, ori_w}, 0.0, 0.0, cv::INTER_CUBIC);

    printf("Save result.\n");

    int pos = input_image.rfind(".");

    input_image.insert(pos, "_depth");

    cv::imwrite(input_image, depth_map);

    // 清理资源  

    context->destroy();  

    engine->destroy();  

    runtime->destroy();

    delete[] data; data = NULL;

    // delete[] DepthMap[0]; DepthMap[0] = NULL;

 

    return 0;  

}

/**

 * @brief 计算已知坐标的离心距离

 * @param center 

 * @return float

 */

float CalculateCentrifugalDistance(cv::Mat ori_img, cv::Point2f center)

{

    float center_x = ori_img.cols / 2.0;

    float center_y = ori_img.rows / 2.0;

    return sqrt(pow((center.x-center_x), 2) + pow((center.y-center_y), 2));

}

  • 14
    点赞
  • 23
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值