class MentalAttenStateDataset(Dataset):
    def __init__(self, data_dir, subj="", direction='h'):
        """
        Custom dataset example for reading image locations and labels from folder
        but reading images from files

        Args:
            data_dir (string): path to raw file
            subj (string): subject no
            direction (string): h or v
        """
        self.data_dir = data_dir
        self.direction = direction
        self.subj = subj
        self.label_dict = {'f':0, 'u':1, 'd':2}
        # Get eeg list
        self.eeg_list = glob.glob(data_dir+'\\' + self.subj + '*')
        # Calculate len
        self.data_len = len(self.eeg_list)

    def __getitem__(self, index):
        filepath = self.eeg_list[index]
        filename = filepath.split('\\')[-1]
        data = np.load(filepath)
        if self.direction == 'h':
            data = data.T
        data = np.expand_dims(data, axis = 0)
        label = self.label_dict.get(filename.split('_')[2])
        return data, label

    def __len__(self):
        return self.data_len