Spaces:
Running
Running
File size: 4,283 Bytes
3aa4060 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import os
import random
import numpy as np
import torch
from grad.utils import fix_len_compatibility
from grad_extend.utils import parse_filelist
class TextMelSpeakerDataset(torch.utils.data.Dataset):
def __init__(self, filelist_path):
super().__init__()
self.filelist = parse_filelist(filelist_path, split_char='|')
self._filter()
print(f'----------{len(self.filelist)}----------')
def _filter(self):
items_new = []
# segment = 200
items_min = 250 # 10ms * 250 = 2.5 S
items_max = 500 # 10ms * 400 = 5.0 S
for mel, vec, pit, spk in self.filelist:
if not os.path.isfile(mel):
continue
if not os.path.isfile(vec):
continue
if not os.path.isfile(pit):
continue
if not os.path.isfile(spk):
continue
temp = np.load(pit)
usel = int(temp.shape[0] - 1) # useful length
if (usel < items_min):
continue
if (usel >= items_max):
usel = items_max
items_new.append([mel, vec, pit, spk, usel])
self.filelist = items_new
def get_triplet(self, item):
# print(item)
mel = item[0]
vec = item[1]
pit = item[2]
spk = item[3]
use = item[4]
mel = torch.load(mel)
vec = np.load(vec)
vec = np.repeat(vec, 2, 0) # 320 VEC -> 160 * 2
pit = np.load(pit)
spk = np.load(spk)
vec = torch.FloatTensor(vec)
pit = torch.FloatTensor(pit)
spk = torch.FloatTensor(spk)
vec = vec + torch.randn_like(vec) # Perturbation
len_vec = vec.size()[0] - 2 # for safe
len_pit = pit.size()[0]
len_min = min(len_pit, len_vec)
mel = mel[:, :len_min]
vec = vec[:len_min, :]
pit = pit[:len_min]
if len_min > use:
max_frame_start = vec.size(0) - use - 1
frame_start = random.randint(0, max_frame_start)
frame_end = frame_start + use
mel = mel[:, frame_start:frame_end]
vec = vec[frame_start:frame_end, :]
pit = pit[frame_start:frame_end]
# print(mel.shape)
# print(vec.shape)
# print(pit.shape)
# print(spk.shape)
return (mel, vec, pit, spk)
def __getitem__(self, index):
mel, vec, pit, spk = self.get_triplet(self.filelist[index])
item = {'mel': mel, 'vec': vec, 'pit': pit, 'spk': spk}
return item
def __len__(self):
return len(self.filelist)
def sample_test_batch(self, size):
idx = np.random.choice(range(len(self)), size=size, replace=False)
test_batch = []
for index in idx:
test_batch.append(self.__getitem__(index))
return test_batch
class TextMelSpeakerBatchCollate(object):
# mel: [freq, length]
# vec: [len, 256]
# pit: [len]
# spk: [256]
def __call__(self, batch):
B = len(batch)
mel_max_length = max([item['mel'].shape[-1] for item in batch])
max_length = fix_len_compatibility(mel_max_length)
d_mel = batch[0]['mel'].shape[0]
d_vec = batch[0]['vec'].shape[1]
d_spk = batch[0]['spk'].shape[0]
# print("d_mel", d_mel)
# print("d_vec", d_vec)
# print("d_spk", d_spk)
mel = torch.zeros((B, d_mel, max_length), dtype=torch.float32)
vec = torch.zeros((B, max_length, d_vec), dtype=torch.float32)
pit = torch.zeros((B, max_length), dtype=torch.float32)
spk = torch.zeros((B, d_spk), dtype=torch.float32)
lengths = torch.LongTensor(B)
for i, item in enumerate(batch):
y_, x_, p_, s_ = item['mel'], item['vec'], item['pit'], item['spk']
mel[i, :, :y_.shape[1]] = y_
vec[i, :x_.shape[0], :] = x_
pit[i, :p_.shape[0]] = p_
spk[i] = s_
lengths[i] = y_.shape[1]
# print("lengths", lengths.shape)
# print("vec", vec.shape)
# print("pit", pit.shape)
# print("spk", spk.shape)
# print("mel", mel.shape)
return {'lengths': lengths, 'vec': vec, 'pit': pit, 'spk': spk, 'mel': mel}
|