深度估计

深度估计数据集用于训练模型，以近似图像中每个像素相对于相机的距离，即深度。这些数据集所支持的应用主要集中在视觉机器感知和机器人感知等领域。典型应用包括为自动驾驶汽车绘制街道地图。本指南将展示如何对深度估计数据集应用变换。

在开始之前，请确保已安装最新版本的 albumentations：

pip install -U albumentations

Albumentations 是一个用于计算机视觉数据增强的 Python 库。它支持多种计算机视觉任务，例如图像分类、目标检测、分割和关键点估计。

本指南使用 NYU Depth V2 数据集，该数据集由来自各种室内场景的视频序列组成，通过 RGB 相机和深度相机录制。该数据集包含来自 3 个城市的场景，并提供图像及其深度图作为标签。

加载数据集的 train 分割，并查看一个示例：

>>> from datasets import load_dataset

>>> train_dataset = load_dataset("sayakpaul/nyu_depth_v2", split="train")
>>> index = 17
>>> example = train_dataset[index]
>>> example
{'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=640x480>,
 'depth_map': <PIL.TiffImagePlugin.TiffImageFile image mode=F size=640x480>}

数据集包含两个字段：

image：一个具有 uint8 数据类型的 PIL PNG 图像对象。
depth_map: 一个 PIL Tiff 图像对象，其数据类型为 float32，表示该图像的深度图。

在此处，深度图采用 TIFF 格式，因为它支持广泛的数据类型，包括 float32 数据。但值得注意的是，JPEG/PNG 格式只能存储 uint8 或 uint16 数据。因此，如果您的深度图以 JPEG/PNG 格式保存，请使用 Image(mode="F") 类型将其加载为单通道 float32，如同普通深度图一样：

>>> from datasets import Image

>>> train_dataset = train_dataset.cast_column("depth_map", Image(mode="F"))

接下来，使用以下命令查看图像：

>>> example["image"]

在查看深度图之前，我们需要先使用 .convert('RGB') 将其数据类型转换为 uint8，因为 PIL 无法显示 float32 图像。现在来看看它对应的深度图：

>>> example["depth_map"].convert("RGB")

全是黑色！你需要为深度图添加一些颜色才能正确可视化。为此，我们可以在显示时使用 plt.imshow() 自动应用颜色，或者使用 plt.cm 创建彩色深度图然后再显示。本示例采用了后者，因为之后可以保存/写入彩色深度图。（下面的实用函数取自 FastDepth 仓库）。

>>> import numpy as np
>>> import matplotlib.pyplot as plt

>>> cmap = plt.cm.viridis

>>> def colored_depthmap(depth, d_min=None, d_max=None):
...     if d_min is None:
...         d_min = np.min(depth)
...     if d_max is None:
...         d_max = np.max(depth)
...     depth_relative = (depth - d_min) / (d_max - d_min)
...     return 255 * cmap(depth_relative)[:,:,:3]

>>> def show_depthmap(depth_map):
...    if not isinstance(depth_map, np.ndarray):
...        depth_map = np.array(depth_map)
...    if depth_map.ndim == 3:
...        depth_map = depth_map.squeeze()

...    d_min = np.min(depth_map)
...    d_max = np.max(depth_map)
...    depth_map = colored_depthmap(depth_map, d_min, d_max)

...    plt.imshow(depth_map.astype("uint8"))
...    plt.axis("off")
...    plt.show()

>>> show_depthmap(example["depth_map"])

您还可以可视化多张不同的图像及其对应的深度图。

>>> def merge_into_row(input_image, depth_target):
...     if not isinstance(input_image, np.ndarray):
...         input_image = np.array(input_image)
...
...     d_min = np.min(depth_target)
...     d_max = np.max(depth_target)
...     depth_target_col = colored_depthmap(depth_target, d_min, d_max)
...     img_merge = np.hstack([input_image, depth_target_col])
...
...     return img_merge

>>> random_indices = np.random.choice(len(train_dataset), 9).tolist()
>>> plt.figure(figsize=(15, 6))
>>> for i, idx in enumerate(random_indices):
...     example = train_dataset[idx]
...     ax = plt.subplot(3, 3, i + 1)
...     image_viz = merge_into_row(
...         example["image"], example["depth_map"]
...     )
...     plt.imshow(image_viz.astype("uint8"))
...     plt.axis("off")

现在使用 albumentations 应用一些数据增强。增强变换包括：

随机水平翻转
随机裁剪
随机亮度和对比度
随机伽马校正
随机色相饱和度

>>> import albumentations as A

>>> crop_size = (448, 576)
>>> transforms = [
...     A.HorizontalFlip(p=0.5),
...     A.RandomCrop(crop_size[0], crop_size[1]),
...     A.RandomBrightnessContrast(),
...     A.RandomGamma(),
...     A.HueSaturationValue()
... ]

此外，定义一个映射以更好地反映目标键名。

>>> additional_targets = {"depth": "mask"}
>>> aug = A.Compose(transforms=transforms, additional_targets=additional_targets)

定义 additional_targets 后，你可以将目标深度图传递给 aug 的 depth 参数，而不是 mask。你将在下方定义的 apply_transforms() 函数中注意到这一更改。

创建一个函数，将变换应用于图像及其深度图：

>>> def apply_transforms(examples):
...     transformed_images, transformed_maps = [], []
...     for image, depth_map in zip(examples["image"], examples["depth_map"]):
...         image, depth_map = np.array(image), np.array(depth_map)
...         transformed = aug(image=image, depth=depth_map)
...         transformed_images.append(transformed["image"])
...         transformed_maps.append(transformed["depth"])
...
...     examples["pixel_values"] = transformed_images
...     examples["labels"] = transformed_maps
...     return examples

使用 set_transform() 函数对数据集的批次进行即时转换，以减少磁盘空间占用：

>>> train_dataset.set_transform(apply_transforms)

您可以通过对示例图像的 pixel_values 和 labels 进行索引来验证转换是否成功：

>>> example = train_dataset[index]

>>> plt.imshow(example["pixel_values"])
>>> plt.axis("off")
>>> plt.show()

在图像对应的深度图上可视化相同的变换：

>>> show_depthmap(example["labels"])

您也可以复用之前的 random_indices 来可视化多个训练样本：

>>> plt.figure(figsize=(15, 6))

>>> for i, idx in enumerate(random_indices):
...     ax = plt.subplot(3, 3, i + 1)
...     example = train_dataset[idx]
...     image_viz = merge_into_row(
...         example["pixel_values"], example["labels"]
...     )
...     plt.imshow(image_viz.astype("uint8"))
...     plt.axis("off")

Update on GitHub

Datasets

深度估计