Make Your First GAN Using PyTorch/

...

/

The Dataset Class

The Dataset Class

Learn about the dataset class.

We'll cover the following...

To work with images from the CelebA dataset we’ll need to modify the Dataset class that we used with the MNIST images.

The CelebADataset class

The changes are fairly simple because we’ve already seen how to open a HDF5 file and extract images from it. Have a look at this CelebADataset class.

Press + to interact
class CelebADataset(Dataset):
    
    def __init__(self, file):
        self.file_object = h5py.File(file, 'r')
        self.dataset = self.file_object['img_align_celeba']
        pass
    
    def __len__(self):
        return len(self.dataset)
    
    def __getitem__(self, index):
        if (index >= len(self.dataset)):
          raise IndexError()
        img = numpy.array(self.dataset[str(index)+'.jpg'])
        return torch.FloatTensor(img) / 255.0
    
    def plot_image(self, index):
        plt.imshow(numpy.array(self.dataset[str(index)+'.jpg']), interpolation='nearest')
        pass
    
    pass

  • The constructor __init__() opens the HDF5 file and also opens the group named img_align_celeba, ready for individual ...