一、简单介绍
单目深度估计(MDE)正受到越来越多的关注,这得益于它在广泛的下游任务中的基础性作用。精确的深度信息不仅在经典应用中如3D重建、导航和自动驾驶中受到青睐,在现代场景中如AI生成的内容,包括图像、视频和3D场景中也备受关注。
Depth Anything v2在精细程度、建模、透明物体反射物体识别上相比v1和其它算法都有进一步的提升,v2与v1相比在算法上没有主要改进,但是在训练数据集上有较大改动,详情可以参阅官方论文。
二、ONNX推理
1.首先从社区中下载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)
四、总结
1.在rknn模型转换时一定要注意均值方差的设置。
2.rknn不支持动态输入,需要进行固定。
3.用虚拟机效率比较低,除了没有用npu加速外,模型算子里有一个很大的转置没有使用到npu,一直在用cpu跑。所以感觉在板子上速度也很慢。