探索深度学习中的稳定扩散和仿射不变深度估计

5 阅读26分钟

在计算机视觉领域,深度学习技术已经成为解决许多任务的主要方法,其中包括图像分类、目标检测、语义分割等。然而,在单目深度估计方面,仍然存在着一些挑战。单目深度估计是指从单张RGB图像中估计场景中每个像素点的深度信息,这是许多视觉任务的关键前提。本文介绍了两种前沿的方法:稳定扩散和仿射不变深度估计,旨在解决深度估计中的关键问题。在本节中,我们将简要介绍这两个概念,并探讨研究的背景和动机。

1.1 稳定扩散简介

稳定扩散是一种基于概率的生成模型,其核心思想是利用扩散过程和自回归建模来生成清晰的图像。扩散过程模拟了噪声图像逐步转变为清晰图像的过程,通过在每个步骤中逐渐添加噪声,生成高质量的图像。稳定扩散的引入为图像生成任务带来了新的思路和方法,为深度估计提供了新的可能性。在稳定扩散中,图像被视为由初始噪声图像逐步演化而来,每一步都引入了更多的噪声,直到生成清晰的图像。这种逐步扩散的过程利用了概率分布的性质,通过不断迭代生成更高质量的图像。稳定扩散的工作原理和概念使得我们能够从噪声图像中准确地生成清晰的图像,为深度估计任务提供了新的思路和方法。

1.2 仿射不变深度估计简介

仿射不变深度估计是一种利用深度学习方法从单目图像中估计场景深度信息的技术。与传统的深度估计方法不同,仿射不变深度估计具有更强的鲁棒性,能够在不受姿态变化和光照变化影响的情况下准确地估计深度。这使得该方法在实际场景中具有更广泛的适用性和稳健性。通过使用深度学习方法,特别是卷积神经网络(CNN),我们可以学习到图像中的特征和场景结构,并从中推断出深度信息。仿射不变深度估计的方法和技术使得我们能够在各种复杂的场景中准确地估计深度,为实际应用提供了强大的工具和方法。

1.3 研究背景和动机

在当前的计算机视觉领域,单目深度估计一直是一个备受关注的研究方向。然而,传统的深度估计方法往往受限于姿态变化、光照变化等因素的影响,难以实现精确的深度估计。因此,引入了稳定扩散和仿射不变深度估计的概念,旨在解决这一挑战,并为单目深度估计领域带来新的突破。通过结合稳定扩散的生成能力和仿射不变深度估计的鲁棒性,我们期望能够提出一种更加有效和准确的深度估计方法,为视觉任务的实现和发展做出贡献。

在接下来的部分中,我们将详细介绍稳定扩散和仿射不变深度估计的原理、方法和实验结果,以及对未来研究方向的展望。通过深入研究和探索,我们希望能够为深度估计领域的进步和发展提供新的思路和方法。

2. 稳定扩散和深度估计

稳定扩散的工作原理是一种基于概率的生成模型,其核心思想是通过扩散过程和自回归建模来逐步生成图像。在稳定扩散中,图像被视为由初始噪声图像逐步演化而来,每个时间步都会引入一定程度的噪声,直到生成清晰的图像。这个过程类似于在图像上逐步叠加细节和结构,直到最终图像呈现出所需的特征和内容。

dt=αtd0+1αtεd_t=\sqrt{{\overline\alpha}}_td_0+\sqrt{1-\overline{\alpha_t}}\varepsilon
L=Ed0,ϵN(0,I),tU(T)ϵϵ^22ℒ = \mathbb{E}_{d_0, \epsilon \sim \mathcal{N}(0, I), t \sim U(T)} \left\| \epsilon - \hat{\epsilon} \right\|_2^2

具体而言,稳定扩散模型的生成过程可以分为以下几个步骤:

  1. 初始化阶段:首先,随机生成一个噪声图像作为初始输入。这个噪声图像可以是随机像素值的矩阵,也可以是根据某种分布生成的初始图像。
  2. 扰动过程:接下来,在每个时间步中,对当前图像进行扰动操作。这个扰动操作通常由一个高斯噪声所参数化,即在当前图像的基础上加上一个服从高斯分布的随机噪声。这个噪声的强度可以随着时间步的增加而逐渐减小,这意味着随着时间的推移,图像会逐渐变得更加清晰。
  3. 逐步减小扰动强度:在每个时间步中,扰动的强度会逐渐减小。这意味着随着时间的推移,噪声对图像的影响会逐渐减弱,最终图像会逐渐收敛到一个稳定状态,呈现出清晰的结构和内容。

通过这种逐步的扰动和减小噪声的过程,稳定扩散模型能够逐步生成清晰的图像。这种生成过程结合了概率建模和自回归建模的特点,能够捕捉图像的局部和全局结构,并生成具有高质量的图像。稳定扩散模型的这种生成方式使其在图像生成和深度估计等任务中具有广泛的应用前景。

2.2 深度估计的挑战与方法介绍

在深度估计领域,面临着多种挑战,如视角变化、光照变化、遮挡等因素的影响。传统的深度估计方法通常基于几何学原理,利用图像的纹理、边缘等特征进行深度估计。然而,这些方法往往依赖于复杂的数学模型和假设,对于复杂的场景和光照条件不具备良好的鲁棒性。

为了克服这些挑战,研究者们转向了基于深度学习的方法。基于深度学习的深度估计方法通过学习大量的数据中图像特征和深度信息的映射关系来实现深度估计,具有更高的灵活性和鲁棒性。这些方法通常采用卷积神经网络(CNN)来提取图像的特征,并通过回归或分类的方式来推断深度信息。

具体来说,基于深度学习的深度估计方法可以分为以下两类:

  1. 基于传统几何的方法:这类方法通常利用图像的纹理、边缘等特征进行深度估计,如结构光、三角测距等。这些方法依赖于复杂的数学模型和假设,对场景的复杂性和光照条件敏感。
  2. 基于深度学习的方法:这类方法利用深度学习技术从大量数据中学习到图像特征和深度信息的映射关系,实现更准确和鲁棒的深度估计。它们通常采用卷积神经网络(CNN)来提取图像特征,并通过回归或分类的方式来推断深度信息。随着深度学习技术的发展,这些方法在深度估计中取得了巨大的进展。

近年来,随着深度学习技术的不断发展和完善,基于深度学习的深度估计方法不断涌现,并在训练集和测试集上取得了良好的性能。例如,端到端的深度估计方法能够直接从图像中推断深度信息,极大地简化了深度估计的流程。此外,一些基于生成对抗网络(GAN)的方法通过生成逼真的深度图像来提升深度估计的准确性和鲁棒性。

这些基于深度学习的方法为解决深度估计中的挑战提供了新的思路和方法,为实现更准确、更鲁棒的深度估计在实际应用中提供了重要的技术支持。

3. 实现细节

3.1 混合数据集训练

混合数据集训练是一种有效的策略,能够提高深度学习模型在不同数据分布下的泛化能力。在本文中,我们采用了两个合成数据集进行训练:Hypersim和Virtual KITTI。这两个数据集具有不同的分辨率和长宽比,因此能够提供更丰富和多样化的训练样本,有助于模型学习到更广泛的场景和特征。

在混合数据集训练过程中,我们以一定的概率从这两个数据集中随机选择样本,并将它们组合成一个批次。这样的做法能够使模型在训练过程中接触到不同数据分布下的样本,从而更好地适应各种真实场景的变化和复杂性。通过在训练集上引入不同数据集的样本,模型可以学习到更多的特征和变化,提高其对于真实数据的适应能力。

然而,混合数据集训练的关键在于确定数据集之间的混合比例。在本文中,我们通过调整伯努利参数来控制数据集的选择概率。具体来说,我们尝试了不同的混合比例,并在验证集上评估模型的性能。通过对模型在不同混合比例下的性能进行比较,我们能够找到最佳的混合比例,使模型在不同数据分布下均能取得良好的性能表现。

这种混合数据集训练的策略能够帮助模型更好地适应各种真实场景的数据分布,提高其泛化能力和鲁棒性。通过合理地选择混合比例,我们可以使模型更有效地学习到数据的特征和结构,从而在实际应用中取得更好的性能表现。

3.2 退火的多分辨率噪声

d~=(dd2d98d20.5)×2\tilde{d} = (\frac{{d - d^2}}{{d_{98} - d^2}} - 0.5 )\times 2

多分辨率噪声是一种在深度学习领域中广泛应用的数据增强技术,其目的在于提高模型的鲁棒性和泛化能力。在本文中,我们将深入探讨多分辨率噪声的原理、应用及改进,并重点介绍引入退火策略后的改进方法。

  • 多分辨率噪声的原理和应用

多分辨率噪声是一种通过在图像中添加不同分辨率的噪声来进行数据增强的方法。其基本原理是在模型训练过程中,为图像引入不同程度的噪声,以模拟现实世界中的各种噪声环境,并帮助模型更好地学习到鲁棒的特征。通常,多分辨率噪声会在图像的不同尺度上引入高斯噪声或其他类型的噪声,以模拟真实世界中的噪声情况。

在实际应用中,多分辨率噪声通常被用于图像分类、目标检测、语义分割等任务中。通过在训练数据中引入多样化的噪声,可以使模型更好地适应各种复杂的场景和环境,提高其对噪声和干扰的抵抗能力。此外,多分辨率噪声还可以有效地防止模型出现过拟合的情况,提高其在测试数据上的泛化能力。

  • 引入退火策略的改进方法

尽管多分辨率噪声在数据增强中发挥着重要作用,但其性能仍然受到一些限制。为了进一步提高多分辨率噪声的效果,我们引入了退火策略,即在噪声的每个层级上逐步减小噪声的强度。这样的改进能够更好地保留图像中的高质量特征,从而提高模型的性能和泛化能力。

具体而言,我们修改了多分辨率噪声的权重计算公式,引入了退火因子 ( s_i )。这个退火因子随着扩散过程的进行而逐渐减小,使得在噪声消失的时间步上,低分辨率层级的影响变得更小。这样一来,模型在训练过程中能够更好地学习到高质量的特征,从而提高了其性能和泛化能力。

  • 实验结果及效果分析

在进行了改进的多分辨率噪声训练后,我们对模型进行了一系列的实验,并评估了其性能和泛化能力。实验结果表明,引入退火策略后的多分辨率噪声训练在各种任务和数据集上都取得了显著的改进。模型在测试数据集上的性能表现更加稳定,且对噪声和干扰的抵抗能力得到了显著提高。这表明退火策略能够有效地改进多分辨率噪声的效果,并提高模型的性能和泛化能力。

  • 潜在应用和未来展望

改进后的多分辨率噪声训练方法具有广泛的潜在应用和未来展望。首先,它可以应用于各种深度学习任务中,如图像分类、

目标检测、语义分割等,以提高模型的性能和鲁棒性。其次,可以进一步研究和优化退火策略,探索更加有效的权重调整方法,进一步提高多分辨率噪声的效果。最后,可以将改进后的多分辨率噪声训练方法应用于更广泛的领域,如医学影像分析、自动驾驶等,为这些领域的发展提供新的思路和方法。

综上所述,改进后的多分辨率噪声训练方法具有重要的理论意义和实际应用价值,将为深度学习领域的发展和应用提供新的思路和方法。

3.3 与地面真实深度的对齐

mins1,,sNt1,,tN(1bi=1N1j=i+1Nd^id^j22+λR)\min_{s_1, \ldots, s_N \\ t_1, \ldots, t_N} (\sqrt{ \frac{1}{b} \sum_{i=1}^{N-1} \sum_{j=i+1}^{N} \| \hat{d}'_i - \hat{d}'_j \|_2^2} + \lambda R)

在深度估计任务中,评估模型的性能通常需要将预测的深度图与地面真实深度进行对齐。对齐的过程可以确保预测的深度图与地面真实深度图之间有一个合适的映射关系,从而可以进行有效的性能评估。在本文中,我们采用了一种基于最小二乘拟合的方法来进行深度图的对齐,以提高评估的准确性和可靠性。

具体而言,我们首先将预测的深度图和地面真实深度图转换为向量形式,并且只考虑地面真实深度中存在的像素。这样做的目的是为了消除深度图中不可见区域的影响,从而提高对齐的准确性。然后,我们使用最小二乘法来拟合预测的深度值和地面真实深度值之间的线性关系。最小二乘法可以找到一条直线,使得预测的深度值与地面真实深度值之间的残差平方和最小,从而得到一个尺度因子和一个偏移量。

在得到尺度因子和偏移量之后,我们将预测的深度图按照这个尺度因子和偏移量进行缩放和平移,以使其与地面真实深度图尽可能接近。这样一来,预测的深度图就可以与地面真实深度图在尺度和位置上对齐,从而可以进行有效的性能评估。这种基于最小二乘拟合的对齐方法具有简单、有效的特点,可以提高深度估计任务的评估准确性和可靠性,为模型的优化和改进提供重要参考。

3.4 3D可视化

在深度估计任务中,进行3D可视化操作是一种非常重要的手段,可以帮助我们更直观地理解模型的预测结果,并对模型的性能进行评估和优化。在本节中,我们将详细介绍实验结果,并对每个实验部分进行更深入的分析和讨论,以及评估模型性能的影响。

我们计算了预测深度图与地面真实深度图之间的比例和偏移标量。这些计算结果可以帮助我们了解模型预测的深度值与实际深度值之间的偏差情况,从而评估模型的准确性。通过比较比例和偏移标量,我们可以发现模型在不同场景和数据集上的性能差异,并且可以针对性地进行模型调整和优化。

我们使用相机内参将像素投影到度量3D空间中。这一步操作可以将预测的深度图转换为真实世界中的三维点云数据,从而更直观地展示模型对场景的理解和深度估计能力。通过观察三维点云数据,我们可以发现模型在不同区域和物体上的深度估计准确性,以及模型对场景结构的理解程度。

对于没有地面真实深度的“现实世界”样本,我们估计了尺度、偏移和相机内参等参数。这些参数估计结果可以帮助我们将模型的预测结果与真实世界场景进行对比,从而评估模型的泛化能力和鲁棒性。通过比较模型预测的深度图与真实世界场景的差异,我们可以发现模型在不同环境和条件下的性能表现,并且可以提出改进模型的建议和方法。

为了可视化法线,我们对每个位置进行了最小二乘平面拟合,考虑到其周围3×3像素的邻域范围。通过可视化法线,我们可以更直观地展示模型对场景表面的理解和深度估计结果。通过观察法线的方向和变化,我们可以发现模型在不同区域和物体上的深度估计准确性,并且可以评估模型对场景几何结构的理解程度。

通过这些实现细节和操作步骤,我们可以更全面地评估模型的性能和鲁棒性,发现模型存在的问题和不足之处,并提出改进模型的方法和策略。这些实验结果和分析将为我们进一步优化模型,提高深度估计的准确性和鲁棒性提供重要参考和指导。

4. 实验结果

4.1 含深度的稳定扩散VAE

我们对预训练的稳定扩散VAE在深度图生成任务中的表现进行了评估。深度图生成是一个关键的任务,因为准确的深度图对于许多计算机视觉应用如三维重建、虚拟现实等具有重要意义。通过稳定扩散VAE生成的深度图,我们可以评估其对于不同场景和对象的表现。实验结果显示,稳定扩散VAE生成的深度图与地面真实深度之间的平均绝对误差(MAE)相对较小,表明其在深度图生成任务中表现出良好的性能和准确性。

在进一步分析中,我们可以将稳定扩散VAE生成的深度图与其他深度估计方法生成的深度图进行比较。这种比较可以帮助我们评估稳定扩散VAE在深度估计方面的优势和局限性,以及确定其在不同场景和数据集上的通用性和鲁棒性。通过与其他方法的比较,我们可以更全面地了解稳定扩散VAE在深度估计中的表现,并为进一步的研究提供指导和启示。

此外,我们还可以对稳定扩散VAE生成的深度图进行定性分析,以便直观地了解其生成效果。通过可视化稳定扩散VAE生成的深度图和地面真实深度之间的差异和一致性,我们可以深入了解其在不同场景和条件下的表现特点,并从中发现模型的优势和改进空间。

4.2 VAE解码器后通道的一致性

我们进一步评估了稳定扩散潜在空间对深度表示的适用性。这一部分的重点在于分析从VAE解码器获取的深度通道的一致性,并探讨这种一致性与解码器输出的关系。通过对潜在空间解码的深度通道之间的差异进行定量分析,我们可以评估稳定扩散潜在空间对深度表示的稳定性和可靠性。

在进一步分析中,我们可以探讨潜在空间解码的深度通道之间的差异与深度估计精度之间的关系。通过研究这种关系,我们可以深入了解稳定扩散潜在空间对深度表示的影响因素,并为提高深度估计精度提供参考和建议。此外,我们还可以将稳定扩散潜在空间与其他深度表示方法进行比较,以评估其相对优势和适用性。

4.3 预测方差和训练噪声

另一个重要的实验部分是评估模型在不同训练和推断条件下的预测一致性。在这一部分中,我们对三种不同的训练方式进行了比较,包括使用高斯噪声、多分辨率噪声和退火的多分辨率噪声。我们通过对模型进行不同噪声条件下的训练,并在验证集上进行推断,来评估模型的预测一致性和鲁棒性。

进一步分析中,我们可以探讨不同训练噪声条件对模型性能和泛化能力的影响。通过研究不同训练噪声条件下模型的预测一致性和鲁棒性,我们可以深入了解噪声对深度估计模型的影响,并为模型的改进和优化提供指导和建议。

4.4 混合训练数据集的比例

我们研究了在微调协议中使用合成数据集的影响。通过在Hypersim和Virtual KITTI两个合成数据集上进行训练,并在真实数据集上进行验证,我们评估了混合不同数据集比例对模型性能的影响。进一步分析中,我们可以探讨不同数据集比例对模型性能和泛化能力的影响,以及确定最佳的数据集混合比例。

通过对混合训练数据集比例的影响进行深入分析,我们可以了解合成数据集在模型训练中的作用和价值,以及确定合适的数据集比例以提高模型性能和泛化能力。此外,我们还可以探讨不同数据集比例对模型在不同场景和条件下的适用性和鲁棒性的影响,为模型的进一步优化和改进提供指导和建议。


部署过程:

1. 选择检查点

在这个部分,用户可以通过下拉菜单选择要使用的深度估计模型。具体代码如下:

ckpt_dic = {
    "Original (higher quality)": "prs-eth/marigold-v1-0",
    "LCM (faster)": "prs-eth/marigold-lcm-v1-0",
}
​
w = widgets.Dropdown(
    options=['Original (higher quality)', 'LCM (faster)'],
    value=ckpt_name,
    description='Checkpoint:',
)
​
def on_change(change):
    if change['type'] == 'change' and change['name'] == 'value':
        ckpt_name = change['new']
        ckpt_path = ckpt_dic[ckpt_name]
​
w.observe(on_change)
​
display(w)

这段代码定义了一个包含两个选项的字典 ckpt_dic,分别代表两个不同模型的检查点。然后,创建了一个下拉菜单 w,供用户选择模型。通过观察下拉菜单的变化,可以实时更新所选择的检查点。

2. 准备输入数据

在这个部分,用户可以选择使用提供的样本图片或上传自己的图片进行深度估计。具体代码如下:

button_download = widgets.Button(description="Use sample images")
button_upload = widgets.Button(description="⬆ Upload images")
button_clear_in = widgets.Button(description="♻ Clear input folder")
button_clear_out = widgets.Button(description="♻ Clear output folder")
out_box = widgets.Output()
​
# 定义按钮点击时的回调函数,分别用于下载样本图片、上传图片、清空输入文件夹和清空输出文件夹
def on_button_download_clicked(b):
    # 下载样本图片并放置到指定文件夹
    pass
​
def on_button_upload_clicked(b):
    # 上传图片到指定文件夹
    pass
​
def on_button_clear_in_clicked(b):
    # 清空输入文件夹
    pass
​
def on_button_clear_out_clicked(b):
    # 清空输出文件夹
    pass
​
button_download.on_click(on_button_download_clicked)
button_upload.on_click(on_button_upload_clicked)
button_clear_in.on_click(on_button_clear_in_clicked)
button_clear_out.on_click(on_button_clear_out_clicked)
​
# 将按钮放置在适当的布局中
widgets.VBox([widgets.HBox([button_upload, button_download]),
              widgets.HBox([button_clear_in, button_clear_out]),
              out_box])

这段代码创建了四个按钮,分别用于下载样本图片、上传图片、清空输入文件夹和清空输出文件夹。每个按钮的点击事件都会触发相应的回调函数,执行对应的操作。

3. 显示输入图片

在这个部分,展示用户上传或选择的输入图片。具体代码如下:

image_paths = glob(os.path.join(input_dir, "*"))
display_images(image_paths)

通过 glob 函数获取输入文件夹中的所有图片路径,然后使用自定义函数 display_images 将这些图片显示在Notebook中。

4. 运行深度估计

在这个部分,使用所选的模型进行深度估计。具体代码如下:

# 设置深度估计的参数
denoising_steps = 4
ensemble_size = 5
processing_res = 768
match_input_res = True# 执行深度估计
for rgb_path in tqdm(rgb_filename_list, desc=f"Estimating depth", leave=True):
    # 读取输入图片
    input_image = Image.open(rgb_path)
​
    # 预测深度
    pipeline_output = pipe(
        input_image,
        denoising_steps=denoising_steps,
        ensemble_size=ensemble_size,
        processing_res=processing_res,
        match_input_res=match_input_res,
        batch_size=0,
        color_map="Spectral",
        show_progress_bar=True,
    )
​
    # 处理预测结果
    depth_pred: np.ndarray = pipeline_output.depth_np
    depth_colored: Image.Image = pipeline_output.depth_colored
​
    # 保存结果
    # (此处省略了保存结果的代码)

在这段代码中,首先设置了深度估计的参数,然后遍历了所有的输入图片路径,逐张进行深度估计。深度估计的结果会保存在指定的输出文件夹中。

5. 显示预测结果

在这个部分,展示深度估计的预测结果。具体代码如下:

output_images_colored  = glob(os.path.join(output_dir_color, "*"))
display_images(output_images_colored)

使用 glob 函数获取输出文件夹中的所有彩色深度图像路径,然后使用自定义函数 display_images 将这些图像显示在Notebook中。


run.py

深度估计模型部署指南

在本指南中,我们将介绍如何使用Marigold库进行单张图像的深度估计。我们将逐步解释如何设置和运行代码,以及如何解释和优化结果。

1. 准备环境和数据

首先,确保你已经安装了所需的Python库。在本示例中,我们使用了Marigold、torch、PIL和tqdm等库。如果还未安装,你可以使用pip进行安装。

pip install marigold torch torchvision pillow tqdm

2. 下载预训练模型

在运行代码之前,我们需要下载并指定要使用的预训练模型。在本示例中,我们使用了Marigold库中的一个预训练模型,具体路径为prs-eth/marigold-lcm-v1-0

checkpoint_path = "prs-eth/marigold-lcm-v1-0"

3. 设置参数和文件夹路径

在代码中,我们需要指定输入图像文件夹路径和输出结果保存的文件夹路径。此外,我们还可以调整一些深度估计的参数,例如去噪步数、集成预测数量、处理分辨率等。

parser = argparse.ArgumentParser(
    description="Run single-image depth estimation using Marigold."
)
parser.add_argument(
    "--checkpoint",
    type=str,
    default="prs-eth/marigold-lcm-v1-0",
    help="Checkpoint path or hub name.",
)
​
parser.add_argument(
    "--input_rgb_dir",
    type=str,
    required=True,
    help="Path to the input image folder.",
)
​
parser.add_argument(
    "--output_dir", type=str, required=True, help="Output directory."
)
​
# inference setting
parser.add_argument(
    "--denoise_steps",
    type=int,
    default=4,
    help="Diffusion denoising steps, more steps results in higher accuracy but slower inference speed. For the original (DDIM) version, it's recommended to use 10-50 steps, while for LCM 1-4 steps.",
)
parser.add_argument(
    "--ensemble_size",
    type=int,
    default=5,
    help="Number of predictions to be ensembled, more inference gives better results but runs slower.",
)
parser.add_argument(
    "--half_precision",
    "--fp16",
    action="store_true",
    help="Run with half-precision (16-bit float), might lead to suboptimal result.",
)
​
# resolution setting
parser.add_argument(
    "--processing_res",
    type=int,
    default=768,
    help="Maximum resolution of processing. 0 for using input image resolution. Default: 768.",
)
parser.add_argument(
    "--output_processing_res",
    action="store_true",
    help="When input is resized, out put depth at resized operating resolution. Default: False.",
)
parser.add_argument(
    "--resample_method",
    choices=["bilinear", "bicubic", "nearest"],
    default="bilinear",
    help="Resampling method used to resize images and depth predictions. This can be one of `bilinear`, `bicubic` or `nearest`. Default: `bilinear`",
)
​
# depth map colormap
parser.add_argument(
    "--color_map",
    type=str,
    default="Spectral",
    help="Colormap used to render depth predictions.",
)
​
# other settings
parser.add_argument(
    "--seed",
    type=int,
    default=None,
    help="Reproducibility seed. Set to `None` for unseeded inference.",
)
​
parser.add_argument(
    "--batch_size",
    type=int,
    default=0,
    help="Inference batch size. Default: 0 (will be set automatically).",
)
parser.add_argument(
    "--apple_silicon",
    action="store_true",
    help="Flag of running on Apple Silicon.",
)
​
args = parser.parse_args()

4. 运行深度估计

通过解析命令行参数或手动设置参数,我们可以执行深度估计。代码会遍历输入图像文件夹中的所有图像,并对每张图像进行深度估计。深度估计结果将保存为.npy和.png格式的文件,并根据深度值着色保存为彩色图片。

with torch.no_grad():
    os.makedirs(output_dir, exist_ok=True)

    for rgb_path in tqdm(rgb_filename_list, desc="Estimating depth", leave=True):
        # Read input image
        input_image = Image.open(rgb_path)

        # Predict depth
        pipe_out = pipe(
            input_image,
            denoising_steps=denoise_steps,
            ensemble_size=ensemble_size,
            processing_res=processing_res,
            match_input_res=match_input_res,
            batch_size=batch_size,
            color_map=color_map,
            show_progress_bar=True,
            resample_method=resample_method

,
            seed=seed,
        )

        depth_pred: np.ndarray = pipe_out.depth_np
        depth_colored: Image.Image = pipe_out.depth_colored

        # Save as npy
        rgb_name_base = os.path.splitext(os.path.basename(rgb_path))[0]
        pred_name_base = rgb_name_base + "_pred"
        npy_save_path = os.path.join(output_dir_npy, f"{pred_name_base}.npy")
        if os.path.exists(npy_save_path):
            logging.warning(f"Existing file: '{npy_save_path}' will be overwritten")
        np.save(npy_save_path, depth_pred)

        # Save as 16-bit uint png
        depth_to_save = (depth_pred * 65535.0).astype(np.uint16)
        png_save_path = os.path.join(output_dir_tif, f"{pred_name_base}.png")
        if os.path.exists(png_save_path):
            logging.warning(f"Existing file: '{png_save_path}' will be overwritten")
        Image.fromarray(depth_to_save).save(png_save_path, mode="I;16")

        # Colorize
        colored_save_path = os.path.join(
            output_dir_color, f"{pred_name_base}_colored.png"
        )
        if os.path.exists(colored_save_path):
            logging.warning(
                f"Existing file: '{colored_save_path}' will be overwritten"
            )
        depth_colored.save(colored_save_path)

5. 结果分析和优化

在获得深度估计结果后,我们可以通过可视化和分析结果来评估模型的性能和鲁棒性。我们还可以尝试调整模型的参数,例如调整去噪步数、集成预测数量等,以进一步优化模型的性能和效果。

通过以上步骤,我们可以轻松地部署和使用Marigold库进行单张图像的深度估计,并且可以根据需要对模型进行进一步的优化和调整。