data_loader_vcc2020.py

from torch.utils import data
from sklearn.preprocessing import StandardScaler
import torch
import glob
import os
from os.path import join, basename, dirname, split, exists
import numpy as np
import h5py

def read_hdf5(hdf5_name, hdf5_path='feats'):
    """Read hdf5 dataset.

    Args:
        hdf5_name (str): Filename of hdf5 file.
        hdf5_path (str): Dataset name in hdf5 file.

    Return:
        any: Dataset values.

    """
    if not os.path.exists(hdf5_name):
        raise Exception(f"There is no such a hdf5 file ({hdf5_name}).")
        sys.exit(1)

    hdf5_file = h5py.File(hdf5_name, "r")

    if hdf5_path not in hdf5_file:
        raise Exception(f"There is no such a data in hdf5 file. ({hdf5_path})")
        sys.exit(1)

    hdf5_data = hdf5_file[hdf5_path][()]
    hdf5_file.close()

    return hdf5_data

def to_categorical(y, num_classes=None):
    """Converts a class vector (integers) to binary class matrix.
    E.g. for use with categorical_crossentropy.
    # Arguments
        y: class vector to be converted into a matrix
            (integers from 0 to num_classes).
        num_classes: total number of classes.
    # Returns
        A binary matrix representation of the input. The classes axis
        is placed last.
    From Keras np_utils
    """
    y = np.array(y, dtype='int')
    input_shape = y.shape
    if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:
        input_shape = tuple(input_shape[:-1])
    y = y.ravel()
    if not num_classes:
        num_classes = np.max(y) + 1
    n = y.shape[0]
    categorical = np.zeros((n, num_classes), dtype=np.float32)
    categorical[np.arange(n), y] = 1
    output_shape = input_shape + (num_classes,)
    categorical = np.reshape(categorical, output_shape)
    return categorical


class MyDataset(data.Dataset):
    """Dataset for MCEP features and speaker labels."""

    def __init__(self, speakers, data_dir, min_length = 128):
        self.speakers = speakers
        self.min_length = min_length
        self.mc_files = []
        self.spk2files = {}
        self.scaler = None
        
        for spk in speakers:
            if spk not in self.spk2files:
                self.spk2files[spk] = []
            if exists(f"{data_dir}/{spk}_raw/{spk}_train"):
                files = glob.glob(f"{data_dir}/{spk}_raw/{spk}_train/*.h5")
                print(spk)
                print(len(files))
                self.spk2files[spk].extend(files)
                for f in files:
                    self.mc_files.append((spk, f))
        self.spk2idx = dict(
            zip(self.speakers, range(len(self.speakers))))  

        
        print(f"loading files {len(self.mc_files)}")

        self.num_files = len(self.mc_files)
        

    def rm_too_short_utt(self, mc_files, min_length):
        new_mc_files = []
        for spk,mc_file in mc_files:
            mc = np.load(mc_file)
            if mc.shape[0] > min_length:
                new_mc_files.append((spk, mc_file))
        return new_mc_files

    def sample_seg(self, feat):
        #assert feat.shape[0] - self.min_length >= 0
        if feat.shape[0] < self.min_length:
            feat = np.pad(feat, [[0,self.min_length - feat.shape[0]],[0,0]])
        s = np.random.randint(0, feat.shape[0] - self.min_length + 1)
        if self.scaler is not None:
        
            return self.scaler.transform(feat[s:s + self.min_length, :])
        else:
            return feat[s:s + self.min_length, :]
    def __len__(self):
        return self.num_files

    def __getitem__(self, index):
        src_spk, src_filename = self.mc_files[index]
        
        
        if src_spk not in self.speakers:
            raise Exception(f"speaker {src_spk} not in self.speakers {self.speakers}")
        src_spk_idx = self.spk2idx[src_spk]

        #src_mc = np.load(src_filename)
        src_mc = read_hdf5(src_filename)
        src_mc = self.sample_seg(src_mc)   
        src_mc = np.transpose(src_mc, (1, 0))  
        # to one-hot
        src_spk_cat = np.squeeze(to_categorical([src_spk_idx], num_classes=len(self.speakers)))
        
        # sample target speaker, source speaker is excluded
        speakers = self.speakers[:]
        speakers.remove(src_spk)
        
        trg_spk_sample = np.random.randint(0, len(speakers))
        trg_spk = speakers[trg_spk_sample]
        trg_spk_idx = self.speakers.index(trg_spk)
        trg_spk_cat = np.squeeze(to_categorical([trg_spk_idx], num_classes = len(self.speakers)))
        # sample one target speaker feature file, will be the input to the speaker encoder
        trg_spk_files = self.spk2files[trg_spk]
        trg_file_sample = np.random.randint(0, len(trg_spk_files))
        trg_filename = trg_spk_files[trg_file_sample]

        #trg_mc = np.load(trg_filename)
        trg_mc = read_hdf5(trg_filename)
        # segment length also min_length
        trg_mc = self.sample_seg(trg_mc)
        trg_mc = np.transpose(trg_mc, (1,0))


        return torch.FloatTensor(src_mc), torch.LongTensor([src_spk_idx]).squeeze_(), torch.FloatTensor(src_spk_cat), torch.FloatTensor(trg_mc), torch.LongTensor([trg_spk_idx]).squeeze_(), torch.FloatTensor(trg_spk_cat)


class TestDataset(object):

    def __init__(self, speakers, data_dir,  src_spk, trg_spk):
        self.speakers = speakers 
        self.data_dir = data_dir
        self.spk2idx = dict(zip(self.speakers, range(len(self.speakers))))
        self.prefix_length = len(self.speakers[0])
        self.src_spk = src_spk
        self.trg_spk = trg_spk
        mc_files = []
        mc_files.extend(sorted(glob.glob(join(data_dir, src_spk+'_raw',src_spk+'_test', '*.h5'))))
        self.mc_files = mc_files
        if len(self.mc_files) == 0:
            raise Exception(f"found no mc files in path {data_dir}")


        self.src_mel_dir = f'{data_dir}/{src_spk}'

        self.spk_idx_src, self.spk_idx_trg = self.spk2idx[src_spk], self.spk2idx[trg_spk]
        spk_cat_src = to_categorical([self.spk_idx_src], num_classes=len(self.speakers))
        spk_cat_trg = to_categorical([self.spk_idx_trg], num_classes=len(self.speakers))
        self.spk_c_org = spk_cat_src
        self.spk_c_trg = spk_cat_trg

    def get_batch_test_data(self, batch_size=8, sample_id = None):
        batch_data = []
        if sample_id is not None:
            print(f"sample_id {sample_id}")
            trg_file = self.trg_mc_files[0]
            for s_id in sample_id:
                mc_file_name = join(self.ft_data_dir, self.src_spk, f'{self.src_spk}_{s_id}-feats.npy')
                batch_data.append((mc_file_name, trg_file))
                
        else:    
            for i in range(batch_size):
                mc_file = self.mc_files[i]  # ./data/dump/train_nodev/norm/SSB0005/SSB00050084-feats.npy
                if self.use_sp_enc:
                    trg_file = self.trg_mc_files[0]
                    batch_data.append( (mc_file, trg_file))  
                else:
                    batch_data.append(mc_file)
        return batch_data





def get_loader(speakers, data_dir, min_length, batch_size=8, mode='train', num_workers=4, ):
    dataset = MyDataset(speakers, data_dir, min_length =min_length)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batch_size,
                                  shuffle=(mode == 'train'),
                                  num_workers=num_workers,
                                  drop_last=True)
    return data_loader