DepthAnything(1): 先跑一跑Depth Anything_depth anything离线怎么跑-CSDN博客
DepthAnything(2): 基于ONNXRuntime在ARM(aarch64)平台部署DepthAnything-CSDN博客
目录
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));
}