Depth Anything v2 ONNX推理以及RKNN推理

い不靠譜︶朱Sir

于 2024-08-15 11:08:48 发布

阅读量376

点赞数 4

分类专栏：实用项目部署文章标签：人工智能计算机视觉深度学习

本文链接：https://blog.csdn.net/qq_45195912/article/details/141216367

版权

实用项目部署专栏收录该内容

2 篇文章 1 订阅

订阅专栏

一、简单介绍

单目深度估计（MDE）正受到越来越多的关注，这得益于它在广泛的下游任务中的基础性作用。精确的深度信息不仅在经典应用中如3D重建、导航和自动驾驶中受到青睐，在现代场景中如AI生成的内容，包括图像、视频和3D场景中也备受关注。

Depth Anything v2在精细程度、建模、透明物体反射物体识别上相比v1和其它算法都有进一步的提升，v2与v1相比在算法上没有主要改进，但是在训练数据集上有较大改动，详情可以参阅官方论文。

二、ONNX推理

1.首先从社区中下载ONNX模型

ONNX模型链接

官方直接给出了ONNX模型，我们就不再根据其文档再从pt转换到ONNX了，建议下载

depth_anything_v2_vits_dynamic.onnx，以方便后续转换为RKNN模型。

2.数据预处理

通过cv库导入图片，确定ONNX模型的地址，对图片进行预处理。

image_path='IMG_1215.HEIC.JPG'
ONNX_MODEL='depth_anything_v2_vits_dynamic.onnx'
image = cv2.imread(str(image_path))
h, w = image.shape[:2]
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
image = cv2.resize(image, (518,518), interpolation=cv2.INTER_CUBIC)
image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
image = image.transpose(2, 0, 1)[None].astype("float32")

3.ONNX推理

使用onnxruntime进行模型推理，这一步你可以选择使用cuda或者cpu。

sess_options = ort.SessionOptions()
sess_options.enable_profiling = False
providers = ["CPUExecutionProvider"]
session = ort.InferenceSession(
    ONNX_MODEL, sess_options=sess_options, providers=providers
)
binding = session.io_binding()
ort_input = session.get_inputs()[0].name
binding.bind_cpu_input(ort_input, image)
ort_output = session.get_outputs()[0].name
binding.bind_output(ort_output, "cpu")
session.run_with_iobinding(binding)

4.对结果进行后处理

这里要注意，出来的结果一定要转成np数组，通道转换，最后使用cv库将图片放大到原大小，这里会出现细节损失。

depth = binding.get_outputs()
depth=np.array(depth)
depth = binding.get_outputs()[0].numpy()

depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
depth = depth.transpose(1, 2, 0).astype("uint8")
depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_CUBIC)
depth_color = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)

cv2.imshow("depth", depth_color)
cv2.waitKey(0)

三、RKNN推理

1.数据预处理

过程与ONNX一样，注意这里使用的是模拟器，如果要在板子上运行需要改一下。

import cv2
import numpy as np
from rknn.api import *

ONNX_MODEL='depth_anything_v2_vits.onnx'
RKNN_MODEL='Correct_depth_anything_v2_vits.rknn'
IMG_PATH='IMG_1215.HEIC.JPG'

image = cv2.imread(str(IMG_PATH))
h, w = image.shape[:2]
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
image = cv2.resize(image, (518, 518), interpolation=cv2.INTER_CUBIC)
image = (image - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
image = image.transpose(2, 0, 1)[None].astype("float32")

2.RKNN推理

注意这里在config要将随机输入进行固定，这里我固定成了(1,3,518,518)。

rknn = RKNN(verbose=True)
rknn.config(mean_values=[0,0,0], std_values=[1,1,1], target_platform='rk3588',dynamic_input=[[[1,3,518,518]]])
ret = rknn.load_onnx(model=ONNX_MODEL)
if ret != 0:
    print('Load model failed!')
print('done')
ret = rknn.build(do_quantization=False)
if ret != 0:
    print('Load model failed!')
print('done')
ret = rknn.init_runtime()
# ret = rknn.init_runtime('rk3566')
if ret != 0:
    print('Init runtime environment failed!')
# print('done')
print('--> Export rknn model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print('Export rknn model failed!')
    # exit(ret)
print('done')
outputs_rknn = rknn.inference(inputs=[image], data_format=['nchw'])

3.对结果进行处理

总的来说与ONNX相似，整个过程没有什么难度。

depth = outputs_rknn[0]
depth=np.array(depth)
depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
depth = depth.transpose(1, 2, 0).astype("uint8")
depth_color = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
cv2.imshow("depth", depth_color)
cv2.waitKey(0)