1d_ecg4_dataset.py

import os
import pandas as pd
from torch.utils.data import Dataset


class PhysioNetDataset(Dataset):
    def __init__(self, ref_file, data_dir):
        self.ref_file = ref_file
        self.data_dir = data_dir
        self.ref_frame = pd.read_csv(ref_file, names=['mat', 'label'])
        self.ref_frame['label'] = pd.Categorical(self.ref_frame['label'])
        self.ref_frame['label_code'] = self.ref_frame['label'].cat.codes

    def __len__(self):
        return len(self.ref_frame)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        mat_name = os.path.join(self.data_dir, self.ref_frame.iloc[idx, 0])

        # extend and create spectrogram
        mat_val = zero_pad(sio.loadmat(mat_name)['val'][0], length=max_length*freq)
        sx = spectrogram(np.expand_dims(mat_val, axis=0))[2]

        # normalize the spectrogram
        sx_norm = (sx - np.mean(sx)) / np.std(sx)

        sample = {'sx': sx_norm, 'label': self.ref_frame.iloc[idx, 2]}

        return sample