文章目录
最近在看深度估计方面的文章,发现了Marigold,它是一种基于扩散模型的单目深度估计技术,通过在Stable Diffusion图像生成模型的基础上进行微调,实现了在多个真实世界数据集上的高准确度和强泛化能力,无需从头开始训练即可高效地预测单个图像的深度信息。接下来我将展示从零开始Marigold的复现。
系统环境
首先介绍一下我所使用的环境,我是在Ubuntu20.04.5,RTX3090(24G显存)的服务器上进行实验的,其中我的磁盘空间为4T(为了下载原始的数据集,后续处理完之后数据集就压缩到很小了,如果你有500G的空间就可以复现该项目,如果你从头开始复现的话,那么你至少要准备一块2T以上的磁盘)。
克隆代码
从论文的摘要部分点击链接进入Github中Marigold的项目地址,或者点击链接进入官方地址,然后使用以下命令来克隆项目到本地或服务器(需要准备提前安装git):
git clone https://github.com/prs-eth/Marigold
cd Marigold
然后使用使用cd命令进入到Marigold项目根目录。
安装conda
环境
根据官方文档的说明,Marigold项目需要在Python 3.10,CUDA 11.7及以上的环境中可以正常运行,这里我使用的环境是Python3.10,CUDA 11.7。
当把项目克隆到本地之后就会发现在项目的根目录下有environment.yaml
、requirements.txt
、requirements+.txt
、requirements++.txt
四个文件,其中environment.yaml
和requirements.txt
是基础环境依赖项,requirements+.txt
和requirements++.txt
是在项目运行环境已经创建成功之后需要额外下载的依赖,根据官方的说明,这边有两种环境创建方法:
-
使用Mamba命令创建环境(前提是你已经安装Mamba环境,在你的base环境中使用
pip install mamba
命令安装mamba
到你的base环境中然后在你项目的根目录下在终端使用以下命令创建marigold环境)mamba env create -n marigold --file environment.yaml # 根据environment.yaml文件创建环境 conda activate marigold # 激活marigold环境
-
使用
conda
或者pip
创建环境:使用conda
环境创建一个虚拟环境或者使用pip
创建一个Python本机虚拟环境并在其中安装依赖项# 1.使用conda创建,使用requirements.txt文件下载相关依赖项 conda create -n marigold python=3.10 conda activate marigold conda install -r requirements.txt # 2.使用pip创建,使用requirements.txt文件下载相关依赖项 python -m venv venv/marigold source venv/marigold/bin/activate pip install -r requirements.txt
这里需要注意的是如果你使用第二种方法进行创建环境的话,那么如果你使用conda
就尽量一直使用conda
,如果你使用pip
那么就尽量一直使用pip
进行下载其他依赖项,可以使用conda
创建环境然后使用pip
安装依赖,但是不要多次交叉使用,因为conda
和pip
对依赖项的管理是不同的,如果两个交叉使用的话可能会报错(比如说使用conda
下载了一个numpy
1.18版本的,但是pip
又下载了一个numpy
1.25版本的,那么就会产生冲突,后续可能会出现一些奇怪的bug。)我这边使用的是conda
创建环境,环境创建成功之后就一直使用pip
命令。
当你创建好环境并把以上基础依赖项安装成功之后,那么可以使用以下命令安装额外的依赖项,这样就不用以后运行项目的时候再回过头配置环境安装依赖项了
pip install -r requirements+.txt -r requirements.txt
数据集与权重文件
当我们把项目的环境配置完成之后接下来就是重头戏了,根据论文中的说明,Marigold使用训练数据集是Hypersim
和vkitti
两个合成的数据集,vkitti
数据集的RGB数据和Depth数据总共约为14G,但是Hypersim
数据局总共有1.9T,所以在刚开始才说为什么磁盘空间需要至少2T以上。
Hypersim
数据集的下载
Hypersim
是一个带有标签的拥有74K高分辨率HDR计算机生成的逼真的室内场景图像集,用于整体室内场景理解的真实感合成数据集。点击链接到Hypersim
的官方项目下,然后使用命令git clone https://github.com/apple/ml-hypersim
克隆到本地(可以和Marigold同属与一个父目录),然后使用cd ml-hypersim
命令进入到ml-hypersim
项目的根目录下,使用命令 python code/python/tools/dataset_download_images.py --downloads_dir ./datasets
通过脚本下载数据集,这里说明一下由于Hypersim
过于庞大,所以分成了若干个压缩文件,我们需要下载完所有数据集之后再进行解压缩以得到完整的Hypersim
数据集(461个压缩文件)。数据集下载过程如图所示:
以下为数据集的部分截图(由于网络原因,整个下载过程持续4天左右)。
当我们把数据集下载完全之后需要把数据集解压缩,由于文件太多,所以我自己写了一个解压缩脚本,通过脚本来自动识别压缩文件进行解压缩,使用以下命令:python script/unzip_and_delete.py
将Hypersim
数据集解压缩得到文件夹(大概20个小时)。然后在这里下载场景分割文件。
当把数据集和场景分割文件都下载成功之后使用如下命令来预处理Hypersim
数据集(约13个小时):
python script/dataset_preprocess/hypersim/preprocess_hypersim.py --split_csv /path/to/metadata_images_split_scene_v1.csv --dataset_dir /path/to/ml-hypersim/datasets --output_dir ./data/Hypersim/processed# 换成你自己保存的文件地址
# 将处理后的数据集使用以下命令进行打包为tar
cd data/Hypersim/processed/train
tar -cf ../../hypersim_processed_train.tar .
使用同样的方法将val数据集打包为tar,至此,Hypersim
数据集处理结束(100多个G),将以上处理完的数据集放到Marigold的根目录下data
文件夹中。
vkitti
数据集的下载
vkitti
数据集是一个用于自动驾驶和计算机视觉研究的大型视频数据集,它包含了从车辆上采集的同步视频、激光雷达和立体相机数据。提供了丰富的交通场景和动态环境信息,用于评估和开发视觉设计
可以通过链接下载vkitti
的RGB数据集和Depth数据集,然后将压缩包放在Marigold的根目录下data
文件夹中。
评估数据集的下载
由于Marigold在训练的过程中使用评估数据集进行验证训练效果,所以如果我们想要复现该项目的话,需要在开始训练之前下载评估数据集。
Diode
数据集是一个专为室内场景环境设计的计算机视觉和机器人导航研究的深度图像数据集,它通过激光扫描获取了多个室内环境的高分辨率深度图和相应的二维图像,包含了一系列详细的室内场景和深度图像、彩色图像、相机姿态和地面真实深度图,用于三维重建、物体识别和场景理解等视觉任务。
Eth3d
数据集是一个为三维计算机视觉任务设计的高分辨率数据集,它包含了多种室内和室外场景的彩色图像、深度图、表面法线和语义分割图,以及相应的相机姿态信息。该数据集特别强调多视角和多传感器数据的融合,适用于3D重建、物体识别、场景理解等研究领域,旨在推动三维视觉算法的发展和评估。
KITTI
数据集是一个广泛用于自动驾驶和计算机视觉研究的基准数据集,它提供了从移动车辆上采集的丰富传感器数据,包括立体图像、激光雷达点云、GPS/IMU数据以及高精度的同步标注信息。KITTI以其大规模、高质量和多样性而闻名,支持包括物体检测、3D重建、光流估计、视觉里程计和场景理解等多种视觉任务的研究和算法评估。
NYU v2
数据集是一个室内场景理解的数据集,它提供了从多个角度拍摄的高分辨率彩色图像和相应的深度图像,以及精确的像素级标注,包括表面法线和语义分割。该数据集广泛应用于室内场景的三维重建、物体识别和深度估计等计算机视觉任务,因其丰富的标注和多样的视角而成为研究者的重要资源。
Scannet
数据集是一个大规模的室内3D理解数据集,它包含了150000多个室内场景的高分辨率RGB-D图像和对应的3D重建,以及详细的场景标注和物体实例分割。这个数据集为计算机视觉中的3D场景理解、物体识别和室内布局分析等任务提供了丰富的资源,是研究者在室内场景深度学习领域的重要工具。
当所有的数据集下载完成之后的目录结构为:
测试数据的下载
根据官方的说明,使用以下命令下载论文中的测试用例:
bash script/download_sample_data.sh
下载之后将图像存放在input/in-the-wild_example
目录下。
权重文件的下载
marigold-v1-0
和marigold-lcm-v1-0
官方提供了两个已经训练好的权重文件分别为marigold-v1-0
和marigold-lcm-v1-0
,其中LCM权重是从原来的权重文件中提炼出来的,目的是加快推理速度(通过减少推理步骤)。推理可以减少1到4步(默认为1),marigold是与论文中一直的权重文件,如果为了进行学术比较,可以使用marigold。以上两个权重文件的下载方式以及推理运行方式如下(将以上两个模型存放在根目录下的prs-eth
文件夹下):
bash script/download_weights.sh marigold-v1-0
python run.py \
--checkpoint prs-eth/marigold-v1-0 \
--denoise_steps 50 \
--ensemble_size 10 \
--input_rgb_dir input/in-the-wild_example \
--output_dir output/in-the-wild_example
# or LCM checkpoint
bash script/download_weights.sh marigold-lcm-v1-0
python run.py \
--input_rgb_dir input/in-the-wild_example \
--output_dir output/in-the-wild_example_lcm
Stable-diffusion-2
Marigold是从Stable-diffusion-2的模型中微调得到的,所以如果想要从头复现或者训练自己的模型需要从该链接中下载Stable-diffusion-2模型,然后将该模型存放在根目录的base_ckpt_dir
文件夹下
以上权重文件都下载完成之后的项目结构为:
在以上准备工作都完成之后整体的项目文件结构目录如下:
注:这里我用的是VScode连接的服务器,如果不知道怎么使用VScode连接服务器的可以看这篇文章。
推理与训练
推理
根据官方的说明,默认设置对最佳结果进行了优化。但是使用者可以自定义代码的行为:
- 准确度和速度之间的权衡(对于这两个选项,值越大,准确度越高,但推理速度越慢。)
--ensemble_size
:集成中的推理次数。对于 LCM 来说,ensemble_size
比 更重要denoise_steps
。默认值:105 (用于 LCM)。--denoise_steps
:每次推理的去噪步骤数。对于原始版本 (DDIM),建议使用 10-50 步,而对于 LCM 则建议使用 1-4 步。未分配时 (None
),将从模型配置中读取默认设置。默认:10 4 (用于 LCM)None
。
- 默认情况下,推理脚本会将输入图像调整为处理分辨率,然后将预测调整回原始分辨率。这样可以获得最佳质量,因为 Marigold 所衍生的 Stable Diffusion 在 768x768 分辨率下表现最佳。
--processing_res
:处理分辨率;设置为 0 则直接处理输入分辨率。未分配时(None
),将从模型配置中读取默认设置。默认值:768None
。--output_processing_res
:以处理分辨率产生输出,而不是将其上采样到输入分辨率。默认值:False。--resample_method
:用于调整图像和深度预测大小的重采样方法。可以是bilinear
、bicubic
或之一nearest
。默认值:bilinear
。
--half_precision
或者--fp16
:以半精度(16 位浮点数)运行以获得更快的速度并减少 VRAM 使用率,但可能会导致次优结果。--seed
:可以设置随机种子以确保额外的可重复性。默认值:无(无种子)。注意:强制--batch_size 1
有助于提高可重复性。要确保完全可重复性,需要使用确定性模式。--batch_size
:重复推理的批次大小。默认值:0(自动确定最佳值)。--color_map
:用于对深度预测进行着色的颜色图None
。默认值:光谱。设置为可跳过彩色深度图生成。--apple_silicon
:使用Apple Silicon MPS加速。
可以根据自己的需要进行相应的设置,然后使用以下命令进行推理:
python run.py \
--input_rgb_dir input/in-the-wild_example \
--output_dir output/in-the-wild_example_lcm
python run.py \
--checkpoint checkpoint/marigold-v1-0 \
--denoise_steps 50 \
--ensemble_size 10 \
--input_rgb_dir input/in-the-wild_example\
--output_dir output/in-the-wild_example
如果在运行的过程中出现以下报错,说明在依赖项安装的时候产生冲突,Numpy
版本不对,那么使用命令pip install numpy==1.25.0
重新安装新版本的Numpy
就好了
运行成功之后的画面如图所示:
训练
根据官方的说明,在准备了环境、数据集和预训练权重文件之后,可以开始训练模型了,在训练模型之前官方给出提示设置数据目录的环境参数:
export BASE_DATA_DIR=YOUR_DATA_DIR # 训练数据集的目录
export BASE_CKPT_DIR=YOUR_CHECKPOINT_DIR # 预训练权重的目录
这里的意思是把BASE_DATA_DIR
和BASE_CKPT_DIR
放到系统的环境变量里面了,程序开始运行的时候首先从系统环境变量里面去查找数据集和预训练权重文件,如果你不想做这一步也是可以运行训练文件的,只需要在启动训练的时候指定以上两个内容即可。
当一切准备工作就绪,使用以下命令开始运行train.py
文件
python train.py --config config/train_marigold.yaml --output_dir output/marigold --base_data_dir data --base_ckpt_dir base_ckpt_dir/ --add_datetime_prefix
正常情况下你的项目就成功运行了,如果给出报错wandb
找不到信息,那么打开浏览器搜索wandb.ai
注册一个账号,然后在个人信息里面找到一个登录ID,使用命令wandb login yourID
登录即可,或者参考这篇文章。
如果在程序运行过程中,不幸异常中断了,那么可以使用如下命令从检查点恢复:
python train.py --resume_run output/marigold_base/checkpoint/latest --output_dir output/marigold --base_data_dir data --base_ckpt_dir base_ckpt_dir/ --add_datetime_prefix
正常运行的情况下是以下的界面:
当看到如下界面,祝贺你已经成功开始跑训练了,接下来就是等它训练结束了!