Spaces:
Runtime error
Runtime error
File size: 3,243 Bytes
0d80816 |
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 |
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import random
import torch
from torch.nn.utils.rnn import pad_sequence
from utils.data_utils import *
from models.base.base_dataset import (
BaseCollator,
BaseDataset,
BaseTestDataset,
BaseTestCollator,
)
import librosa
class AutoencoderKLDataset(BaseDataset):
def __init__(self, cfg, dataset, is_valid=False):
BaseDataset.__init__(self, cfg, dataset, is_valid=is_valid)
cfg = self.cfg
# utt2melspec
if cfg.preprocess.use_melspec:
self.utt2melspec_path = {}
for utt_info in self.metadata:
dataset = utt_info["Dataset"]
uid = utt_info["Uid"]
utt = "{}_{}".format(dataset, uid)
self.utt2melspec_path[utt] = os.path.join(
cfg.preprocess.processed_dir,
dataset,
cfg.preprocess.melspec_dir,
uid + ".npy",
)
# utt2wav
if cfg.preprocess.use_wav:
self.utt2wav_path = {}
for utt_info in self.metadata:
dataset = utt_info["Dataset"]
uid = utt_info["Uid"]
utt = "{}_{}".format(dataset, uid)
self.utt2wav_path[utt] = os.path.join(
cfg.preprocess.processed_dir,
dataset,
cfg.preprocess.wav_dir,
uid + ".wav",
)
def __getitem__(self, index):
# melspec: (n_mels, T)
# wav: (T,)
single_feature = BaseDataset.__getitem__(self, index)
utt_info = self.metadata[index]
dataset = utt_info["Dataset"]
uid = utt_info["Uid"]
utt = "{}_{}".format(dataset, uid)
if self.cfg.preprocess.use_melspec:
single_feature["melspec"] = np.load(self.utt2melspec_path[utt])
if self.cfg.preprocess.use_wav:
wav, sr = librosa.load(
self.utt2wav_path[utt], sr=16000
) # hard coding for 16KHz...
single_feature["wav"] = wav
return single_feature
def __len__(self):
return len(self.metadata)
def __len__(self):
return len(self.metadata)
class AutoencoderKLCollator(BaseCollator):
def __init__(self, cfg):
BaseCollator.__init__(self, cfg)
def __call__(self, batch):
# mel: (B, n_mels, T)
# wav (option): (B, T)
packed_batch_features = dict()
for key in batch[0].keys():
if key == "melspec":
packed_batch_features["melspec"] = torch.from_numpy(
np.array([b["melspec"][:, :624] for b in batch])
)
if key == "wav":
values = [torch.from_numpy(b[key]) for b in batch]
packed_batch_features[key] = pad_sequence(
values, batch_first=True, padding_value=0
)
return packed_batch_features
class AutoencoderKLTestDataset(BaseTestDataset):
...
class AutoencoderKLTestCollator(BaseTestCollator):
...
|