dataset.py

import torch
from torch.utils.data.dataset import Dataset
import numpy as np


class MuSeDataset(Dataset):
    def __init__(self, data, partition):
        super(MuSeDataset, self).__init__()
        self.partition = partition
        features, labels_valence, labels_arousal = data[partition]['feature'], data[partition]['label_valence'], data[partition]['label_arousal']
        metas = data[partition]['meta']
        self.feature_dim = features[0].shape[-1]
        self.n_samples = len(features)

        feature_lens = []
        label_lens = []
        for feature in features:
            feature_lens.append(len(feature))

        max_feature_len = np.max(np.array(feature_lens))


        self.feature_lens = torch.tensor(feature_lens)

        self.features = [torch.tensor(f, dtype=torch.float) for f in features]

        self.labels_valence = [torch.tensor(l, dtype=torch.float) for l in labels_valence]
        self.labels_arousal = [torch.tensor(l, dtype=torch.float) for l in labels_arousal]

        self.metas = metas
        self.metas = [m.astype(object).tolist() for m in self.metas]
        pass

    def get_feature_dim(self):
        return self.feature_dim

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        '''

        :param idx:
        :return: feature, feature_len, label, meta with
            feature: tensor of shape seq_len, feature_dim
            feature_len: int tensor, length of the feature tensor before padding
            label: tensor of corresponding label(s) (shape 1 for n-to-1, else (seq_len,1))
            meta: list of lists containing corresponding meta data
        '''

        feature = self.features[idx]
        feature_len = self.feature_lens[idx]
        label_valence = self.labels_valence[idx]
        label_arousal = self.labels_arousal[idx]
        meta = self.metas[idx]

        sample = feature, feature_len, label_valence,label_arousal, meta
        return sample