Depth Anything试玩

0. 引言

最近看到一个叫Depth Anything单目深度估计模型火了，于是抽空把代码和模型下载下来体验了一下，发现确实是很强大。

论文链接：https://arxiv.org/pdf/2401.10891.pdf

代码链接：https://github.com/LiheYoung/Depth-Anything

项目主页: https://depth-anything.github.io/

本文只记录如何使用官方代码跑demo，以及如何导出onnx模型并用onnxruntime做部署。

1. 使用官方代码跑demo

首先从GitHub把代码clone下来：

          
              

            git clone https://github.com/LiheYoung/Depth-Anything
          
        

然后安装依赖库：

            
cd Depth-Anything
            
pip install -r requirements.txt
        

依赖库比较少，搭建环境非常简单，装好以后就可以跑demo了。如果是输入图像，那么可以执行下面的脚本：

          
              

            python run.py --encoder <vits | vitb | vitl> --img-path <img-directory | single-img | txt-file> --outdir <outdir>
          
        

比如：

          
              

            python run.py --encoder vitb --img-path ../test.jpg --outdir output/
          
        

--img-path参数可以是单张图像的路径、存放多张图片的文件夹的路径、存放一系列图像路径的TXT文件的路径。目前官方发布了三个模型：depth_anything_vits14.pth,depth_anything_vitb14.pth,depth_anything_vitl14.pth，分别与参数里的vits, vitb,vitl对应。执行上面的命令后，会自动从Huggingface的网站上下载对应的模型。 「不过需要注意的是，国内目前无法访问Huggingface」 。怎么办呢？不用慌，我们可以使用Huggingface的镜像网站。首先在命令行执行下面的命令设置一下环境变量：

          
              

            export HF\_ENDPOINT=https://hf-mirror.com
          
        

然后再运行run.py脚本就可以愉快地跑模型了，脚本会把结果输出到--outdir参数指定的目录中。下面是我用nuScenes数据集中的图片跑出来的结果：

如果需要跑视频，那么可以用run_video.py脚本:

          
              

            python run\_video.py --encoder vitb --video-path assets/examples\_video --outdir output/
          
        

2. onnx模型部署

2.1 导出onnx模型

导出onnx模型的方法可以参考下面这个仓库：

https://github.com/fabio-sim/Depth-Anything-ONNX

把代码下载下来后运行export.py脚本即可导出onnx模型:

          
              

            python export.py --model s # s对应vits模型
          
        

导出的onnx模型默认存放在weights/目录下。在这个脚本中，同样会去Huggingface网站上下载PyTorch模型，因此需要用镜像网站替换一下。替换方法很简单，就是把代码中指定的链接将huggingface.co直接替换为hf-mirror.com。

            
depth_anything.to(device).load_state_dict(
            
      torch.hub.load_state_dict_from_url(
            
          f"https://hf-mirror.com/spaces/LiheYoung/Depth-Anything/resolve/main/checkpoints/depth_anything_vit{model}14.pth",
            
          map_location="cpu",
            
      ),
            
      strict=True,
            
  )
        

另外，这个脚本导出onnx模型的时候是使用的动态参数模型，如果不想用动态参数可以把dynamic_axes参数去掉改成静态模式。导出的onnx模型还可以使用onnx-simplifier工具简化一下。

2.2 用onnxruntime部署onnx模型

之前写了两篇用onnxruntime部署onnx模型的文章，大家可以参考一下：

AI模型部署 | onnxruntime部署YOLOv8分割模型详细教程

AI模型部署 | onnxruntime部署RT-DETR目标检测模型

部署Depth Anything模型也是差不多的流程。加载好onnx模型后，首先需要对输入图像做预处理，预处理的时候需要做减均值再除以标准差对图像数据做规范化，其他处理操作与YOLOv8和RT-DETR是一样的。预处理函数preprocess的实现如下：

            
def preprocess(
            
    bgr_image,
            
    width,
            
    height,
            
    mean=[123.675, 116.28, 103.53],
            
    std=[58.395, 57.12, 57.375],
            
):
            
    image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
            
    image = cv2.resize(image, (width, height)).astype(np.float32)
            
    image = (image - mean) / std
            
    image = np.transpose(image, (2, 0, 1)).astype(np.float32)
            
    input_tensor = np.expand_dims(image, axis=0)
            
    return input_tensor
        

输入数据准备好以后，就可以送入模型进行推理：

          
              

            outputs = session.run(None, {session.get\_inputs()[0].name: input\_tensor})
          
        

得到模型推理结果后，只需要做一点简单的后处理操作就可以了:

            
depth = outputs[0][0]  
            
depth = np.transpose(depth, [1, 2, 0]) #chw->hwc
            

            
depth = cv2.normalize(depth,None, 0,255,cv2.NORM_MINMAX,cv2.CV_8UC1)
            
colormap = cv2.applyColorMap(depth,cv2.COLORMAP_INFERNO)
            
colormap = cv2.resize(colormap,(image_width,image_height))
            
combined_results = cv2.hconcat([image, colormap])
        

后处理的时候首先调用OpenCV的normalize函数将深度图的像素值调整到[0,255]的范围内，然后调用applyColorMap函数对深度图进行伪彩色化（单纯为了好看，可以不做）。最终得到的结果大概是这个样子的：

模型输入尺寸设置为518x518，batch size设置为1，在GeForce RTX 3090显卡上3个模型的耗时如下：

模型	模型精度	耗时
depth_anything_vits14.onnx	FP32	16 ms
depth_anything_vitb14.onnx	FP32	42 ms
depth_anything_vitl14.onnx	FP32	90 ms

4. 参考资料

• https://github.com/LiheYoung/Depth-Anything

• https://github.com/fabio-sim/Depth-Anything-ONNX