深度估測

深度估測數據集用於訓練模型，以近似影像中每個像素相對於相機的相對距離，這也稱為深度。這些數據集所實現的應用主要位於視覺機器感知和機器人感知等領域。範例應用包括為自動駕駛汽車繪製街道地圖。本指南將向您展示如何對深度估測數據集進行轉換。

在開始之前，請確保您已安裝最新版本的 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：一個資料型態為 float32 的 PIL Tiff 影像物件，即該影像的深度圖。

這裡的深度圖使用 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 應用一些增強。增強轉換包括

隨機水平翻轉
隨機裁剪
隨機亮度與對比度
隨機 Gamma 校正
隨機色相飽和度

>>> 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() 函數，在資料集批次處理時即時 (on-the-fly) 套用轉換，以節省硬碟空間。

>>> 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")

在 GitHub 上更新

資料集 (Datasets)

深度估測