PyTorch NYUv2 Dataset Class

PyTorch wrapper for the NYUv2 dataset focused on multi-task learning.
Data sources available: RGB, Semantic Segmentation(13), Surface Normals, Depth Images.

Downloads data from:

• RGB Train
• RGB Test
• Segmentation Train
• Segmentation Test
• Surface Normals
• Depth Images

Example

from nyuv2 import NYUv2
from torchvision import transforms

t = transforms.Compose([transforms.RandomCrop(400), transforms.ToTensor()])
NYUv2(root="/somepath/NYUv2", download=True, 
      rgb_transform=t, seg_transform=t, sn_transform=t, depth_transform=t)

Dataset NYUv2
    Number of datapoints: 795
    Split: train
    Root Location: /somepath/NYUv2
    RGB Transforms: Compose(
                        RandomCrop(size=(400, 400), padding=None)
                        ToTensor()
                    )
    Seg Transforms: Compose(
                        RandomCrop(size=(400, 400), padding=None)
                        ToTensor()
                    )
    SN Transforms: Compose(
                       RandomCrop(size=(400, 400), padding=None)
                       ToTensor()
                   )
    Depth Transforms: Compose(
                          RandomCrop(size=(400, 400), padding=None)
                          ToTensor()
                      )

Notes

• Each source has its own transformation pipeline
• Downloads datasets only for tasks where the passed transform is not None
• Do not flip surface normals, as the output would be incorrect without further processing
• Image size is 480x640, however some have a white border which can be removed by cropping 16px from all sides
• Semantic Segmentation Classes: (0) background, (1) bed, (2) books, (3) ceiling, (4) chair, (5) floor, (6) furniture, (7) objects, (8) painting, (9) sofa, (10) table, (11) tv, (12) wall, (13) window

Requirements

h5py: 2.9.0
pillow: 6.2.0
pytorch: 0.4.0
torchvision: 0.4.0