Marigold: 重新利用基于扩散的图像生成器进行单目深度估计

CVPR 2024（口头报告，最佳论文奖候选）

本仓库是论文"重新利用基于扩散的图像生成器进行单目深度估计"的官方实现。

柯炳鑫、 Anton Obukhov、黄圣宇、 Nando Metzger、 Rodrigo Caye Daudt、 Konrad Schindler

我们提出了Marigold，这是一个用于单目深度估计的扩散模型及其相关的微调协议。其核心原理是利用现代生成图像模型中存储的丰富视觉知识。我们的模型源自Stable Diffusion，并使用合成数据进行微调，可以零样本迁移到未见过的数据，提供最先进的单目深度估计结果。

预览图

📢 新闻

2024-05-28：训练代码已发布。 2024-03-23：添加了LCM v1.0以加快推理速度 - 在<a href="https://huggingface.co/spaces/prs-eth/marigold-lcm"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face%20(LCM)-Space-yellow" height="16"></a>上试用 2024-03-04：被CVPR 2024接收。 2023-12-22：为Diffusers贡献了社区管道。 2023-12-19：更新许可证为Apache License, Version 2.0。 2023-12-08：添加了 <a href="https://huggingface.co/spaces/toshas/marigold"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Space-yellow" height="16"></a> - 免费试用您的图像！ 2023-12-05：添加了<a href="https://colab.research.google.com/drive/12G8reD13DdpMie5ZQlaFNo2WCGeNUH-u?usp=sharing"><img src="https://yellow-cdn.veclightyear.com/835a84d5/08c90eb9-0077-4a0e-9fac-09f02c0c49ba.svg" height="16"></a> - 深入了解我们的推理流程！ 2023-12-04：添加了<a href="https://arxiv.org/abs/2312.02145"><img src="https://img.shields.io/badge/arXiv-PDF-b31b1b" height="16"></a> 论文和推理代码（本仓库）。

🚀 使用方法

我们提供了几种与Marigold交互的方式：

我们将Marigold管道集成到了diffusers 🧨中。查看这个diffusers教程中的许多激动人心的使用场景。
这里提供了一个免费的在线交互式演示：<a href="https://huggingface.co/spaces/prs-eth/marigold-lcm"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face%20(LCM)-Space-yellow" height="16"></a>（感谢HF团队提供的GPU资源）
本地运行演示（需要GPU和nvidia-docker2，参见安装指南）：
1. 论文版本：docker run -it -p 7860:7860 --platform=linux/amd64 --gpus all registry.hf.space/toshas-marigold:latest python app.py
2. LCM版本：docker run -it -p 7860:7860 --platform=linux/amd64 --gpus all registry.hf.space/prs-eth-marigold-lcm:latest python app.py
Google Colab上的扩展演示：<a href="https://colab.research.google.com/drive/12G8reD13DdpMie5ZQlaFNo2WCGeNUH-u?usp=sharing"><img src="https://yellow-cdn.veclightyear.com/835a84d5/08c90eb9-0077-4a0e-9fac-09f02c0c49ba.svg" height="16"></a>
如果您只想查看示例，请访问我们的展示页面：<a href="https://marigoldmonodepth.github.io"><img src="https://yellow-cdn.veclightyear.com/835a84d5/7173763c-8f16-4628-99df-18e36bb9aa97.svg" height="16"></a>
最后，下面给出了使用此代码库进行本地开发的说明。

🛠️ 设置

推理代码已在以下环境中测试：

Ubuntu 22.04 LTS, Python 3.10.12, CUDA 11.7, GeForce RTX 3090 (pip, Mamba)
CentOS Linux 7, Python 3.10.4, CUDA 11.7, GeForce RTX 4090 (pip)
Windows 11 22H2, Python 3.10.12, CUDA 12.3, GeForce RTX 3080 (Mamba)
MacOS 14.2, Python 3.10.12, M1 16G (pip)

🪧 Windows用户注意事项

我们建议在WSL2中运行代码：

按照安装指南安装WSL。
按照安装指南为WSL安装CUDA支持。
在/mnt/<驱动器字母>/中找到您的驱动器；更多详情请查看WSL FAQ。导航到您选择的工作目录。

📦 仓库

克隆仓库（需要git）：

git clone https://github.com/prs-eth/Marigold.git
cd Marigold

💻 依赖项

我们提供了几种安装依赖项的方法。

使用Mamba，可以与Miniforge3一起安装。

Windows用户：将Linux版本安装到WSL中。

安装后，首先需要激活Miniforge：source /home/$USER/miniforge3/bin/activate。

创建环境并在其中安装依赖项：
```
mamba env create -n marigold --file environment.yaml
conda activate marigold
```

使用pip： 或者，创建一个Python原生虚拟环境并在其中安装依赖项：

python -m venv venv/marigold
source venv/marigold/bin/activate
pip install -r requirements.txt

在运行推理脚本之前保持环境处于激活状态。重新启动终端会话后再次激活环境。

🏃 在您的图像上进行测试

📷 准备图像

使用我们论文中的精选图像：
```
bash script/download_sample_data.sh
```
或将您的图像放在一个目录中，例如input/in-the-wild_example下，然后运行以下推理命令。

🚀 使用LCM运行推理（更快）

LCM检查点是从我们原始检查点蒸馏而来，以实现更快的推理速度（通过减少推理步骤）。推理步骤可以少至1（默认）到4步。使用默认LCM设置运行：

python run.py \
    --input_rgb_dir input/in-the-wild_example \
    --output_dir output/in-the-wild_example_lcm

🎮 使用DDIM运行推理（论文设置）

此设置对应于我们的论文。对于学术比较，请使用此设置运行。

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

您可以在output/in-the-wild_example中找到所有结果。祝您使用愉快！

⚙️ 推理设置

默认设置针对最佳结果进行了优化。但是，代码的行为可以自定义：

在准确性和速度之间进行权衡（对于两个选项，较大的值会带来更好的准确性，但推理速度会变慢。）
- --ensemble_size：集成中的推理通道数量。对于LCM，ensemble_size比denoise_steps更重要。默认值：10 5（对于LCM）。
- --denoise_steps：每次推理通道的去噪步骤数。对于原始（DDIM）版本，建议使用10-50步，而对于LCM则为1-4步。未指定时（None），将从模型配置中读取默认设置。默认值：~~10 4（对于LCM）~~ None。
默认情况下，推理脚本会将输入图像调整为"处理分辨率"，然后将预测结果调整回原始分辨率。这样可以获得最佳质量，因为 Marigold 源自 Stable Diffusion，在 768x768 分辨率下表现最佳。
--processing_res：处理分辨率；设为 0 则直接处理输入分辨率。未指定时（None），将从模型配置中读取默认设置。默认值：~~768~~ None。
--output_processing_res：以处理分辨率输出结果，而不是将其上采样到输入分辨率。默认值：False。
--resample_method：用于调整图像和深度预测的重采样方法。可以是 bilinear、bicubic 或 nearest 之一。默认值：bilinear。
--half_precision 或 --fp16：使用半精度（16 位浮点数）运行，以获得更快的速度和更少的显存使用，但可能导致次优结果。
--seed：可以设置随机种子以确保额外的可重现性。默认值：None（无种子）。注意：强制使用 --batch_size 1 有助于提高可重现性。要确保完全可重现，需要使用确定性模式。
--batch_size：重复推理的批量大小。默认值：0（自动确定最佳值）。
--color_map：用于给深度预测着色的色图。默认值：Spectral。设为 None 可跳过彩色深度图生成。
--apple_silicon：使用 Apple Silicon MPS 加速。

⬇ 检查点缓存

默认情况下，检查点存储在 Hugging Face 缓存中。 HF_HOME 环境变量定义了其位置，可以被覆盖，例如：

export HF_HOME=$(pwd)/cache

或者，使用以下脚本在本地下载检查点权重：

bash script/download_weights.sh marigold-v1-0
# 或 LCM 检查点
bash script/download_weights.sh marigold-lcm-v1-0

在推理时，指定检查点路径：

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

🦿 在测试数据集上进行评估 <a name="evaluation"></a>

安装额外的依赖：

pip install -r requirements+.txt -r requirements.txt

设置数据目录变量（在评估脚本中也需要）并将评估数据集下载到相应的子文件夹中：

export BASE_DATA_DIR=<YOUR_DATA_DIR>  # 设置目标数据目录

wget -r -np -nH --cut-dirs=4 -R "index.html*" -P ${BASE_DATA_DIR} https://share.phys.ethz.ch/~pf/bingkedata/marigold/evaluation_dataset/

运行推理和评估脚本，例如：

# 运行推理
bash script/eval/11_infer_nyu.sh

# 评估预测结果
bash script/eval/12_eval_nyu.sh

注意：尽管已设置种子，但在不同硬件上结果可能仍略有不同。

🏋️ 训练

基于之前创建的环境，安装扩展要求：

pip install -r requirements++.txt -r requirements+.txt -r requirements.txt

为数据目录设置环境参数：

export BASE_DATA_DIR=YOUR_DATA_DIR  # 训练数据目录
export BASE_CKPT_DIR=YOUR_CHECKPOINT_DIR  # 预训练检查点目录

将 Stable Diffusion v2 检查点下载到 ${BASE_CKPT_DIR}

准备 Hypersim 和 Virtual KITTI 2 数据集并保存到 ${BASE_DATA_DIR}。Hypersim 预处理请参考此 README。

运行训练脚本

python train.py --config config/train_marigold.yaml

从检查点恢复，例如

python train.py --resume_run output/marigold_base/checkpoint/latest

评估结果

训练期间只更新和保存 U-Net。要使用训练结果进行推理，请用 checkpoint 输出文件夹中的 unet 文件夹替换 Marigold 检查点中的相应文件夹。然后参考此部分进行评估。

注意：虽然已设置随机种子，但在不同硬件上训练结果可能略有不同。建议不间断训练。

✏️ 贡献

请参考此说明。

🤔 故障排除

问题	解决方案
(Windows) WSL 上无效的 DOS bash 脚本	运行 `dos2unix <script_name>` 转换脚本格式
(Windows) WSL 上出错：`Could not load library libcudnn_cnn_infer.so.8. Error: libcuda.so: cannot open shared object file: No such file or directory`	运行 `export LD_LIBRARY_PATH=/usr/lib/wsl/lib:$LD_LIBRARY_PATH`

🎓 引用

请引用我们的论文：

@InProceedings{ke2023repurposing,
      title={Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation},
      author={Bingxin Ke and Anton Obukhov and Shengyu Huang and Nando Metzger and Rodrigo Caye Daudt and Konrad Schindler},
      booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
      year={2024}
}