Estimación de profundidad

Los conjuntos de datos de estimación de profundidad se utilizan para entrenar un modelo que aproxime la distancia relativa de cada píxel de una imagen con respecto a la cámara, también conocida como profundidad. Las aplicaciones habilitadas por estos conjuntos de datos se centran principalmente en áreas como la percepción visual por máquina y la percepción en robótica. Entre las aplicaciones de ejemplo se encuentra el mapeo de calles para coches autónomos. Esta guía te mostrará cómo aplicar transformaciones a un conjunto de datos de estimación de profundidad.

Antes de comenzar, asegúrate de tener instaladas las versiones más recientes de albumentations:

pip install -U albumentations 

Albumentations es una biblioteca de Python para realizar aumentación de datos para visión por computadora. Admite diversas tareas de visión por computadora, como clasificación de imágenes, detección de objetos, segmentación y estimación de puntos clave.

Esta guía utiliza el conjunto de datos NYU Depth V2, que está compuesto por secuencias de vídeo de diversas escenas interiores, grabadas con cámaras RGB y de profundidad. El conjunto de datos consta de escenas de 3 ciudades y proporciona imágenes junto con sus mapas de profundidad como etiquetas.

Cargue la división train del conjunto de datos y examine un ejemplo:

>>> 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>}

El conjunto de datos tiene dos campos:

• image: un objeto de imagen PNG de PIL con tipo de datos uint8.
• depth_map: un objeto de imagen PIL Tiff con tipo de datos float32 que es el mapa de profundidad de la imagen.

Aquí los mapas de profundidad utilizan el formato TIFF, ya que admite una amplia gama de tipos de datos, incluidos datos float32. Sin embargo, cabe mencionar que el formato JPEG/PNG solo puede almacenar datos uint8 o uint16. Por lo tanto, si tienes mapas de profundidad guardados como JPEG/PNG, usa el tipo Image(mode="F") para cargarlos como float32 de un solo canal, al igual que los mapas de profundidad normales:

>>> from datasets import Image

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

A continuación, comprueba una imagen con:

>>> example["image"]

Antes de ver el mapa de profundidad, primero debemos convertir su tipo de datos a uint8 usando .convert('RGB'), ya que PIL no puede mostrar imágenes float32. Ahora, observa su mapa de profundidad correspondiente:

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

¡Todo está en negro! Necesitarás añadir algo de color al mapa de profundidad para visualizarlo correctamente. Para hacerlo, podemos aplicar el color automáticamente durante la visualización usando plt.imshow() o crear un mapa de profundidad con color usando plt.cm y luego mostrarlo. En este ejemplo, hemos utilizado la segunda opción, ya que nos permite guardar/escribir el mapa de profundidad con color más adelante. (la utilidad a continuación se toma del repositorio 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"])

También puede visualizar varias imágenes diferentes y sus mapas de profundidad correspondientes.

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

Ahora aplica algunas aumentaciones con albumentations. Las transformaciones de aumentación incluyen:

• Volteo horizontal aleatorio
• Recorte aleatorio
• Brillo y contraste aleatorios
• Corrección gamma aleatoria
• Saturación de tono aleatoria

>>> 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()
... ]

Además, define un mapeo para reflejar mejor el nombre de la clave objetivo.

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

Con additional_targets definido, puedes pasar los mapas de profundidad objetivo al argumento depth de aug en lugar de mask. Notarás este cambio en la función apply_transforms() definida a continuación.

Cree una función para aplicar la transformación a las imágenes y también a sus mapas de profundidad:

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

Use la función set_transform() para aplicar la transformación al vuelo a los lotes del conjunto de datos y consumir menos espacio en disco:

>>> train_dataset.set_transform(apply_transforms)

Puedes verificar que la transformación funcionó indexando en pixel_values y labels de una imagen de ejemplo:

>>> example = train_dataset[index]

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

Visualiza la misma transformación en el mapa de profundidad correspondiente de la imagen:

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

También puede visualizar varias muestras de entrenamiento reutilizando el random_indices anterior:

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