torchaudio testing

This commit is contained in:
Si11ium
2020-12-17 08:02:28 +01:00
parent c6fdaa24aa
commit 62d9eb6e8f
6 changed files with 264 additions and 55 deletions

@ -1,30 +1,40 @@
import time
from pathlib import Path
import numpy as np
import pickle
from torch.utils.data import Dataset
from ml_lib.modules.util import AutoPadToShape
class TorchMelDataset(Dataset):
def __init__(self, identifier, mel_path, segment_len, hop_len, label, padding=0, transform=None):
self.padding = padding
self.path = next(iter(Path(mel_path).glob(f'{identifier}_*')))
self.segment_len = segment_len
self.m, self.n = str(self.path).split('_')[-2:] # get spectrogram dimensions
self.n = int(self.n.split('.', 1)[0]) # remove .npy
self.m, self.n = (int(i) for i in (self.m, self.n))
self.offsets = list(range(0, self.n - segment_len, hop_len))
def __init__(self, mel_path, sub_segment_len, sub_segment_hop_len, label, audio_file_len,
sampling_rate, mel_hop_len, n_mels, transform=None, auto_pad_to_shape=True):
super(TorchMelDataset, self).__init__()
self.sampling_rate = sampling_rate
self.audio_file_len = audio_file_len
self.padding = AutoPadToShape((1, n_mels , sub_segment_len)) if auto_pad_to_shape else None
self.path = Path(mel_path)
self.sub_segment_len = sub_segment_len
self.mel_hop_len = mel_hop_len
self.sub_segment_hop_len = sub_segment_hop_len
self.n = int((self.sampling_rate / self.mel_hop_len) * self.audio_file_len + 1)
self.offsets = list(range(0, self.n - self.sub_segment_len, self.sub_segment_hop_len))
self.label = label
self.transform = transform
def __getitem__(self, item):
while Path(str(self.path).replace(self.path.suffix, '.lock')).exists():
time.sleep(0.01)
with self.path.open('rb') as mel_file:
mel_spec = pickle.load(mel_file, fix_imports=True)
start = self.offsets[item]
mel_spec = np.load(str(self.path), allow_pickle=True)
if self.padding > 0:
mel_spec = np.pad(mel_spec, pad_width=[(0, 0), (self.padding // 2, self.padding // 2)], mode='mean')
snippet = mel_spec[:, start: start + self.segment_len]
snippet = mel_spec[:, : , start: start + self.sub_segment_len]
if self.transform:
snippet = self.transform(snippet)
if self.padding:
snippet = self.padding(snippet)
return snippet, self.label
def __len__(self):
return len(self.offsets)
return len(self.offsets)