import torch
from torch.utils import data
from torchvision import transforms

class Dataset(data.Dataset):
  'Characterizes a dataset for PyTorch'
  def __init__(self, list_IDs, labels, data, transforms=None):
        'Initialization'
        self.labels = labels
        self.list_IDs = list_IDs
        self.data = data
        self.transforms = transforms

  def __len__(self):
        'Denotes the total number of samples'
        return len(self.list_IDs)

  def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        ID = self.list_IDs[index]

        # Load data and get label
        X = self.data[int(ID)]
        X = X.reshape(3, 32, 32).transpose([1, 2, 0])
        if self.transforms is not None:
            X = self.transforms(X)
        y = self.labels[ID]

        return X, y