mmlearn.datasets.nyuv2.NYUv2Dataset

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

Bases: Dataset[Example]

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.

Methods

__getitem__(idx)[source]

Return RGB and depth images at index idx.

Return type:

Example