72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torch.utils.data
|
||
|
|
import torchaudio
|
||
|
|
import torchvision
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from utils.util import opt_get
|
||
|
|
|
||
|
|
|
||
|
|
class PreprocessedMelDataset(torch.utils.data.Dataset):
|
||
|
|
|
||
|
|
def __init__(self, opt):
|
||
|
|
path = opt['path']
|
||
|
|
cache_path = opt['cache_path'] # Will fail when multiple paths specified, must be specified in this case.
|
||
|
|
if os.path.exists(cache_path):
|
||
|
|
self.paths = torch.load(cache_path)
|
||
|
|
else:
|
||
|
|
path = Path(path)
|
||
|
|
self.paths = [str(p) for p in path.rglob("*.npz")]
|
||
|
|
torch.save(self.paths, cache_path)
|
||
|
|
self.pad_to = opt_get(opt, ['pad_to_samples'], 10336)
|
||
|
|
|
||
|
|
def __getitem__(self, index):
|
||
|
|
with np.load(self.paths[index]) as npz_file:
|
||
|
|
mel = torch.tensor(npz_file['arr_0'])
|
||
|
|
assert mel.shape[-1] <= self.pad_to
|
||
|
|
padding_needed = self.pad_to - mel.shape[-1]
|
||
|
|
mask = torch.zeros_like(mel)
|
||
|
|
if padding_needed > 0:
|
||
|
|
mel = F.pad(mel, (0,padding_needed))
|
||
|
|
mask = F.pad(mask, (0,padding_needed), value=1)
|
||
|
|
|
||
|
|
output = {
|
||
|
|
'mel': mel,
|
||
|
|
'mel_lengths': torch.tensor(mel.shape[-1]),
|
||
|
|
'mask': mask,
|
||
|
|
'mask_lengths': torch.tensor(mask.shape[-1]),
|
||
|
|
'path': self.paths[index],
|
||
|
|
}
|
||
|
|
return output
|
||
|
|
|
||
|
|
def __len__(self):
|
||
|
|
return len(self.paths)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
params = {
|
||
|
|
'mode': 'preprocessed_mel',
|
||
|
|
'path': 'Y:\\separated\\large_mels',
|
||
|
|
'cache_path': 'Y:\\separated\\large_mels.pth',
|
||
|
|
'pad_to_samples': 10336,
|
||
|
|
'phase': 'train',
|
||
|
|
'n_workers': 0,
|
||
|
|
'batch_size': 16,
|
||
|
|
}
|
||
|
|
from data import create_dataset, create_dataloader
|
||
|
|
|
||
|
|
ds = create_dataset(params)
|
||
|
|
dl = create_dataloader(ds, params)
|
||
|
|
i = 0
|
||
|
|
for b in tqdm(dl):
|
||
|
|
#pass
|
||
|
|
torchvision.utils.save_image((b['mel'].unsqueeze(1)+1)/2, f'{i}.png')
|
||
|
|
i += 1
|
||
|
|
if i > 20:
|
||
|
|
break
|