mmlearn.datasets.nyuv2

SUN RGB-D dataset.

Functions

depth_normalize(depth_file, min_depth=0.01, max_depth=50)[source]

Load depth file and convert to disparity image.

Parameters:

  • depth_file (str) – Path to the depth file.

  • min_depth (float, default=0.01) – Minimum depth value to clip the depth image.

  • max_depth (int, default=50) – Maximum depth value to clip the depth image.

Returns:

The normalized depth image.

Return type:

torch.Tensor

Classes

NYUv2Dataset

NYUv2 dataset.
class NYUv2Dataset(root_dir, split='train', return_type='disparity', rgb_transform=None, depth_transform=None)[source]

NYUv2 dataset.

Parameters:

  • root_dir (str) – Path to the root directory of the dataset.

  • split ({"train", "test"}, default="train") – Split of the dataset to use.

  • return_type ({"disparity", "image"}, default="disparity") –

    Return type of the depth images.

    • ”disparity”: Return the depth image as disparity map.

    • ”image”: Return the depth image as a 3-channel image.

  • rgb_transform (Callable[[PIL.Image], torch.Tensor], default=None) – A callable that takes in an RGB PIL image and returns a transformed version of the image as a PyTorch tensor.

  • depth_transform (Callable[[PIL.Image], torch.Tensor], default=None) – A callable that takes in a depth PIL image and returns a transformed version of the image as a PyTorch tensor.

Raises:

ImportError – If opencv-python is not installed.

__getitem__(idx)[source]

Return RGB and depth images at index idx.

Return type:

Example

__len__()[source]

Return the length of the dataset.

Return type:

int

depth_normalize(depth_file, min_depth=0.01, max_depth=50)[source]

Load depth file and convert to disparity image.

Parameters:

  • depth_file (str) – Path to the depth file.

  • min_depth (float, default=0.01) – Minimum depth value to clip the depth image.

  • max_depth (int, default=50) – Maximum depth value to clip the depth image.

Returns:

The normalized depth image.

Return type:

torch.Tensor