-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathdata.py
212 lines (168 loc) · 7.79 KB
/
data.py
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
from torchaudio.datasets import LIBRISPEECH
from torch.utils.data import Dataset
import pyroomacoustics as pra
import numpy as np
from typing import Tuple
from torch import Tensor
import torch
import random
from librosa import resample
from librosa.effects import split
from scipy.stats import special_ortho_group
from scipy.io import wavfile
import scipy
import pandas as pd
import noisereduce as nr
from tqdm import tqdm
def remove_silence(signal, top_db=20, frame_length=2048, hop_length=512):
'''
Remove silence from speech signal
'''
signal = signal.squeeze()
clips = split(signal, top_db=top_db,
frame_length=frame_length, hop_length=hop_length)
output = []
for ii in clips:
start, end = ii
output.append(signal[start:end])
return torch.cat(output)
class LibriSpeechLocations(LIBRISPEECH):
'''
Class of LibriSpeech recordings. Each recording is annotated with a speaker location.
'''
def __init__(self, source_locs, mic_locs, split, random_source_pos=False,
xyz_min=None, xyz_max=None):
super().__init__("./", url=split, download=True)
self.source_locs = source_locs
self.mic_locs = mic_locs
self.random_source_pos = random_source_pos
self.xyz_min = xyz_min
self.xyz_max = xyz_max
def __getitem__(self, n: int) -> Tuple[Tensor, int, str, int, int, int, int, float, float, int]:
if self.random_source_pos:
# replace stored source loc with new random position. This can be used
# during training to increase the size of the dataset.
source_loc = np.random.uniform(
low=self.xyz_min, high=self.xyz_max, size=self.source_locs[n].shape)
else:
source_loc = self.source_locs[n]
mic_locs = self.mic_locs[n]
seed = n
waveform, sample_rate, transcript, speaker_id, chapter_id, utterance_number = super().__getitem__(n)
return (waveform, sample_rate, transcript, speaker_id, utterance_number), source_loc, mic_locs, seed
def one_delay(room_dim, fs, t60, mic_locs, signal, source_loc, anechoic_prob=0.0, snr=None):
'''
Simulate signal propagation using pyroomacoustics for a given source location.
'''
p = np.random.rand()
if p < anechoic_prob:
e_absorption = 1.0
max_order = 0
else:
e_absorption, max_order = pra.inverse_sabine(t60, room_dim)
room = pra.ShoeBox(room_dim, fs=fs, materials=pra.Material(
e_absorption), max_order=max_order)
room.add_source(source_loc, signal=signal.squeeze())
room.add_microphone(mic_locs)
c = room.c
num_mics = mic_locs.shape[1]
tdoas = np.zeros((num_mics+1, num_mics+1))
for i, mic1 in enumerate(mic_locs.transpose(1,0)):
tdoas[i+1, 0] = (np.linalg.norm(source_loc - mic1)
- np.linalg.norm(source_loc - source_loc)) * fs / c
tdoas[0, i+1] = -tdoas[i+1, 0]
for j, mic2 in enumerate(mic_locs.transpose(1,0)):
tdoas[i+1, j+1] = (np.linalg.norm(source_loc - mic1)
- np.linalg.norm(source_loc - mic2)) * fs / c
# we do not add noise here, this is done using data augmentation during training
room.simulate(reference_mic=0, snr=snr)
x = room.mic_array.signals
return x, tdoas, room
def pad_sequence(batch):
# Make all tensor in a batch the same length by padding with zeros
batch = [item.t() for item in batch]
batch = torch.nn.utils.rnn.pad_sequence(
batch, batch_first=True, padding_value=0.)
return batch
class DelaySimulatorDataset(Dataset):
'''
Given a batch of LibrispeechLocation samples, simulate signal
propagation from source to the microphone locations.
'''
def __init__(self, location_dataset, room_dim, in_fs, out_fs, N, N_gt, t60, anechoic_prob,
train=True, lower_bound=0.5, upper_bound=1.5, repeats=1, remove_silence=False,
snr=None):
self.location_dataset = location_dataset
self.room_dim = room_dim
self.in_fs = in_fs
self.out_fs = out_fs
self.N = N
self.N_gt = N_gt
self.t60 = t60
self.anechoic_prob = anechoic_prob
self.train = train
self.lower_bound = lower_bound
self.upper_bound = upper_bound
self.repeats = repeats
self.remove_silence = remove_silence
self.snr = snr
def generate_data(self):
tensors, source_locs, mic_locs, tdoas = [], [], [], []
count = 0
with tqdm(total=len(self.location_dataset)*self.repeats) as pbar:
for i in range(self.repeats):
for (waveform, _, _, _, _), source_loc, mic_loc, seed in self.location_dataset:
signal = waveform.squeeze()
# Resample
if self.in_fs != self.out_fs:
signal = torch.Tensor(resample(signal.numpy(), orig_sr=self.in_fs,
target_sr=self.out_fs,
res_type="kaiser_fast"))
if self.remove_silence:
signal = remove_silence(signal, frame_length=self.N)
# use random seed for training, fixed for val/test
# this controls the randomness in sound propagation when simulating the room
if not self.train:
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
# sample random reverberation time
this_t60 = np.random.uniform(
low=self.t60[0], high=self.t60[1])
x, tdoa, _ = one_delay(room_dim=self.room_dim, fs=self.out_fs, t60=this_t60,
mic_locs=mic_loc, signal=signal,
source_loc=source_loc,
anechoic_prob=self.anechoic_prob,
snr=self.snr)
start_idx = int(self.lower_bound * self.out_fs)
end_idx = int(self.upper_bound * self.out_fs)
signal = signal[start_idx:end_idx].unsqueeze(0)
correction = 41 # for some reason the simulation delays with 41 samples (in 16kHz)
x = x[:, start_idx+correction:end_idx+correction]
# add the transmitted signal
x = np.concatenate((signal, x), axis=0)
tensors += [torch.as_tensor(x, dtype=torch.float)]
source_locs += [torch.as_tensor(source_loc)]
mic_locs += [torch.as_tensor(mic_loc)]
tdoas += [torch.as_tensor(tdoa)]
pbar.update(1)
count = count + 1
#if count == 300:
# break
# Group the list of tensors into a batched tensor
self.tensors = pad_sequence(tensors).unsqueeze(1).permute(0, 1, 3, 2)
self.source_locs = torch.stack(source_locs, dim=0).unsqueeze(1)
self.mic_locs = torch.stack(mic_locs, dim=0)
self.tdoas = torch.stack(tdoas, dim=0)
def save_data(self, LOG_DIR, name):
filename = f"{LOG_DIR}/{name}.pt"
torch.save([self.source_locs, self.mic_locs, self.tensors, self.tdoas], filename)
def load_data(self, LOG_DIR, name):
filename = f"{LOG_DIR}/{name}.pt"
self.source_locs, self.mic_locs, self.tensors, self.tdoas = torch.load(filename)
def __len__(self):
return len(self.tensors)
def __getitem__(self, n: int):
source_loc = self.source_locs[n]
mic_locs = self.mic_locs[n]
return self.tensors[n], source_loc, mic_locs, self.tdoas[n]