HybridDepth

混合深度：强大的深度融合
通过利用焦点深度和单图像先验

阿什坎甘吉¹ ·苏航² ·田果¹

¹伍斯特理工学院² Nvidia 研究中心

我们发布了 HybridDepth 的改进版本，现已提供新功能和优化的性能！

这项工作提出了 HybridDepth。 HybridDepth 是一种实用的深度估计解决方案，基于从相机捕获的焦点堆栈图像。这种方法在多个知名数据集（包括 NYU V2、DDFF12 和 ARKitScenes）上的性能优于最先进的模型。

消息

• 2024 年 10 月 30 日：发布了 HybridDepth 的第 2 版，具有改进的性能和预训练权重。

• 2024-10-30 ：对 TorchHub 的集成支持，以方便模型加载和推理。

• 2024-07-25 ：预训练模型的首次发布。

• 2024-07-23 ：GitHub 存储库和 HybridDepth 模型上线。

用法

Colab 笔记本入门文件

使用 Colab 笔记本快速开始使用 HybridDepth。

使用 TorchHub

您可以直接使用 TorchHub 选择预训练模型。

可用的预训练模型：

• HybridDepth_NYU5 ：使用 5 焦点堆栈输入在 NYU Depth V2 数据集上进行预训练，并训练 DFF 分支和细化层。

• HybridDepth_NYU10 ：使用 10 焦点堆栈输入在 NYU Depth V2 数据集上进行预训练，并训练 DFF 分支和细化层。

• HybridDepth_DDFF5 ：使用 5 焦点堆栈输入在 DDFF 数据集上进行预训练。

• HybridDepth_NYU_PretrainedDFV5 ：在使用 DFV 进行预训练之后，仅使用 NYU Depth V2 数据集使用 5 焦点堆栈在细化层上进行预训练。

 model_name = 'HybridDepth_NYU_PretrainedDFV5' #change thismodel = torch.hub.load('cake-lab/HybridDepth', model_name , pretrained=True)model.eval()

本地安装

1. 克隆存储库并安装依赖项：

 git 克隆 https://github.com/cake-lab/HybridDepth.gitcd HybridDepth
conda env create -f 环境.yml
conda 激活混合深度

2. 下载预训练的权重：

从下面的链接下载模型的权重并将其放置在checkpoints目录中：

• HybridDepth_NYU_FocalStack5

• HybridDepth_NYU_FocalStack10

• HybridDepth_DDFF_FocalStack5

• HybridDepth_NYU_PretrainedDFV_FocalStack5

3. 预言

为了进行推理，您可以运行以下代码：

 # 加载模型 checkpointmodel_path = 'checkpoints/NYUBest5.ckpt'model = DepthNetModule.load_from_checkpoint(model_path)model.eval()model = model.to('cuda')

加载模型后，使用以下代码处理输入图像并获取深度图：

注意：目前， prepare_input_image函数仅支持.jpg图像。如果需要支持其他图像格式，请修改该函数。

 from utils.io import prepare_input_imagedata_dir = 'focal stack images directory' # 文件夹中焦点堆栈图像的路径# 加载焦点堆栈图像focal_stack, rgb_img, focus_dist = prepare_input_image(data_dir)# 使用 torch.no_grad() 运行推理: out = model (rgb_img、focal_stack、focus_dist)metric_深度 = out[0].squeeze().cpu().numpy() # 度量深度

评估

请首先从下面的链接下载模型的权重并将其放置在checkpoints目录中：

• HybridDepth_NYU_FocalStack5

• HybridDepth_NYU_FocalStack10

• HybridDepth_DDFF_FocalStack5

• HybridDepth_NYU_PretrainedDFV_FocalStack5

数据集准备

1. NYU Depth V2 ：按照此处提供的说明下载数据集。

2. DDFF12 ：按照此处提供的说明下载数据集。

3. ARKitScenes ：按照此处提供的说明下载数据集。

在configs目录中设置配置文件config.yaml 。 configs目录中提供了每个数据集的预配置文件，您可以在其中指定路径、模型设置和其他超参数。这是一个配置示例：

 data: class_path: dataloader.dataset.NYUDataModule # dataset.py 中数据加载器模块的路径
  init_args:nyuv2_data_root: "path/to/NYUv2" # 特定数据集的路径img_size: [480, 640] # 根据您的DataModule要求进行调整remove_white_border: Truenum_workers: 0 # 如果使用合成数据则设置为0use_labels: Truemodel: invert_depth: True # 设置如果模型输出反转深度ckpt_path，则为 True：检查点/checkpoint.ckpt

在test.sh脚本中指定配置文件：

 python cli_run.py test --config configs/config_file_name.yaml

然后，执行评估：

光盘脚本
sh 评估.sh

---

训练

安装综合 CUDA 包

安装图像合成所需的基于 CUDA 的包：

 python utils/synthetic/gauss_psf/setup.py 安装

这将安装合成图像所需的包。

训练配置

在configs目录中设置配置文件config.yaml ，指定数据集路径、批量大小和其他训练参数。以下是使用 NYUv2 数据集进行训练的示例配置：

模型：invert_深度：True
  # 学习率
  lr: 3e-4 # 根据需要调整
  # 权重衰减
  wd: 0.001 # 根据需要调整 data: class_path: dataloader.dataset.NYUDataModule # dataset.py 中数据加载器模块的路径
  init_args:nyuv2_data_root: "path/to/NYUv2" # 数据集 pathimg_size: [480, 640] # 调整 NYUDataModuleremove_white_border: Truebatch_size: 24 # 根据可用内存进行调整num_workers: 0 # 如果使用合成 datause_labels: Trueckpt_path: null 则设置为 0

在train.sh脚本中指定配置文件：

 python cli_run.py train --config configs/config_file_name.yaml

执行训练命令：

光盘脚本
sh 火车.sh

引文

如果我们的工作对您的研究有帮助，请按如下方式引用：

 @misc{ganj2024hybriddepthrobustmetricdepth，title={HybridDepth：通过利用焦点深度和单图像先验进行鲁棒度量深度融合}，作者={Ashkan Ganj、Hang Su和TianGuo}，year={2024}，eprint={2407.18443} , archivePrefix={arXiv}, PrimaryClass={cs.CV}，url={https://arxiv.org/abs/2407.18443}，
}