forked from mrq/DL-Art-School
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
|
import os
|
||
|
|
import random
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.utils.data
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from data.util import get_image_paths, is_wav_file
|
||
|
|
from models.tacotron2.taco_utils import load_wav_to_torch
|
||
|
|
|
||
|
|
|
||
|
|
class WavfileDataset(torch.utils.data.Dataset):
|
||
|
|
|
||
|
|
def __init__(self, opt):
|
||
|
|
self.path = os.path.dirname(opt['path'])
|
||
|
|
cache_path = os.path.join(self.path, 'cache.pth')
|
||
|
|
if os.path.exists(cache_path):
|
||
|
|
self.audiopaths = torch.load(cache_path)
|
||
|
|
else:
|
||
|
|
print("Building cache..")
|
||
|
|
self.audiopaths = get_image_paths('img', opt['path'], qualifier=is_wav_file)[0]
|
||
|
|
torch.save(self.audiopaths, cache_path)
|
||
|
|
self.max_wav_value = 32768.0
|
||
|
|
self.sampling_rate = 24000
|
||
|
|
self.window = 2 * self.sampling_rate
|
||
|
|
|
||
|
|
def get_audio_for_index(self, index):
|
||
|
|
audiopath = self.audiopaths[index]
|
||
|
|
filename = os.path.join(self.path, audiopath)
|
||
|
|
audio, sampling_rate = load_wav_to_torch(filename)
|
||
|
|
if sampling_rate != self.sampling_rate:
|
||
|
|
raise ValueError(f"Input sampling rate does not match specified rate {self.sampling_rate}")
|
||
|
|
audio_norm = audio / self.max_wav_value
|
||
|
|
audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False)
|
||
|
|
return audio_norm, audiopath
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
clip1, clip2 = None, None
|
||
|
|
|
||
|
|
while clip1 is None and clip2 is None:
|
||
|
|
# Split audio_norm into two tensors of equal size.
|
||
|
|
audio_norm, filename = self.get_audio_for_index(index)
|
||
|
|
if audio_norm.shape[0] < self.window * 2:
|
||
|
|
# Try next index. This adds a bit of bias and ideally we'd filter the dataset rather than do this.
|
||
|
|
index = (index + 1) % len(self)
|
||
|
|
continue
|
||
|
|
j = random.randint(0, audio_norm.shape[0] - self.window)
|
||
|
|
clip1 = audio_norm[j:j+self.window]
|
||
|
|
j = random.randint(0, audio_norm.shape[0]-self.window)
|
||
|
|
clip2 = audio_norm[j:j+self.window]
|
||
|
|
|
||
|
|
return {
|
||
|
|
'clip1': clip1.unsqueeze(0),
|
||
|
|
'clip2': clip2.unsqueeze(0),
|
||
|
|
'path': filename,
|
||
|
|
}
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
return len(self.audiopaths)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
params = {
|
||
|
|
'mode': 'wavfile_clips',
|
||
|
|
'path': 'E:\\audio\\LibriTTS\\train-other-500',
|
||
|
|
'phase': 'train',
|
||
|
|
'n_workers': 0,
|
||
|
|
'batch_size': 16,
|
||
|
|
}
|
||
|
|
from data import create_dataset, create_dataloader, util
|
||
|
|
|
||
|
|
ds, c = create_dataset(params, return_collate=True)
|
||
|
|
dl = create_dataloader(ds, params, collate_fn=c)
|
||
|
|
i = 0
|
||
|
|
m = []
|
||
|
|
max_text = 0
|
||
|
|
max_mel = 0
|
||
|
|
for b in tqdm(dl):
|
||
|
|
pass
|
||
|
|
m=torch.stack(m)
|
||
|
|
print(m.mean(), m.std())
|