Depth estimation

Depth estimation datasets are used to train a model to approximate the relative distance of every pixel in an image from the camera, also known as depth. The applications enabled by these datasets primarily lie in areas like visual machine perception and perception in robotics. Example applications include mapping streets for self-driving cars. This guide will show you how to apply transformations to a depth estimation dataset.

Before you start, make sure you have up-to-date versions of albumentations installed:

pip install -U albumentations

Albumentations is a Python library for performing data augmentation for computer vision. It supports various computer vision tasks such as image classification, object detection, segmentation, and keypoint estimation.

This guide uses the NYU Depth V2 dataset which is comprised of video sequences from various indoor scenes, recorded by RGB and depth cameras. The dataset consists of scenes from 3 cities and provides images along with their depth maps as labels.

Load the train split of the dataset and take a look at an example:

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

The dataset has two fields:

image: a PIL PNG image object with uint8 data type.
depth_map: a PIL Tiff image object with float32 data type which is the depth map of the image.

It is mention-worthy that JPEG/PNG format can only store uint8 or uint16 data. As the depth map is float32 data, it can’t be stored using PNG/JPEG. However, we can save the depth map using TIFF format as it supports a wider range of data types, including float32 data.

Next, check out an image with:

>>> example["image"]

Before we look at the depth map, we need to first convert its data type to uint8 using .convert('RGB') as PIL can’t display float32 images. Now take a look at its corresponding depth map:

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

It’s all black! You’ll need to add some color to the depth map to visualize it properly. To do that, either we can apply color automatically during display using plt.imshow() or create a colored depth map using plt.cm and then display it. In this example, we have used the latter one, as we can save/write the colored depth map later. (the utility below is taken from the FastDepth repository).

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

You can also visualize several different images and their corresponding depth maps.

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

Now apply some augmentations with albumentations. The augmentation transformations include:

Random horizontal flipping
Random cropping
Random brightness and contrast
Random gamma correction
Random hue saturation

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

Additionally, define a mapping to better reflect the target key name.

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

With additional_targets defined, you can pass the target depth maps to the depth argument of aug instead of mask. You’ll notice this change in the apply_transforms() function defined below.

Create a function to apply the transformation to the images as well as their depth maps:

>>> 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 the set_transform() function to apply the transformation on-the-fly to batches of the dataset to consume less disk space:

>>> train_dataset.set_transform(apply_transforms)

You can verify the transformation worked by indexing into the pixel_values and labels of an example image:

>>> example = train_dataset[index]

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

Visualize the same transformation on the image’s corresponding depth map:

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

You can also visualize multiple training samples reusing the previous 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")

Datasets

Depth estimation