mmlearn.datasets.sunrgbd.SUNRGBDDataset

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

Bases: Dataset[Example]

SUN RGB-D 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. If “disparity”, the depth images are converted to disparity similar to the ImageBind implementation. Otherwise, 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.

References

Methods

__getitem__(idx)[source]

Return RGB and depth images at index idx.

Return type:

Example