This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 3.0 United States License. 2 | 3 | # Locate and Beamform: Two-dimensional Locating All-neural Beamformer for Multi-channel Speech Separation (LaBNet) 4 | 5 | Official PyTorch implementation and dataset generation scripts of the Interspeech 2023 paper ["Locate and Beamform: Two-dimensional Locating All-neural Beamformer for Multi-channel Speech Separation"](https://arxiv.org/abs/2305.10821) by Fu Yanjie et al. 6 | 7 | ## How to cite 8 | 1. Fu, Yanjie, et al. "Locate and Beamform: Two-dimensional Locating All-neural Beamformer for Multi-channel Speech Separation." arXiv preprint arXiv:2305.10821 (2023). 9 | ```bibtex 10 | @misc{fu2023locate, 11 | title={Locate and Beamform: Two-dimensional Locating All-neural Beamformer for Multi-channel Speech Separation}, 12 | author={Yanjie Fu and Meng Ge and Honglong Wang and Nan Li and Haoran Yin and Longbiao Wang and Gaoyan Zhang and Jianwu Dang and Chengyun Deng and Fei Wang}, 13 | year={2023}, 14 | eprint={2305.10821}, 15 | archivePrefix={arXiv}, 16 | primaryClass={eess.AS} 17 | } 18 | ``` 19 | 20 | ## Model 21 | We implement LaBNet based on Generalized spatio-temporal RNN Beamformer (GRNN-BF), which directly learns the beamforming weights from the estimated speech and noise spatial covariance matrices. For more details, please refer to the original paper: ["Generalized Spatio-Temporal RNN Beamformer for Target Speech Separation"](https://www.isca-speech.org/archive/pdfs/interspeech_2021/xu21i_interspeech.pdf). 22 | 23 | The figure below shows the architecture of our proposed LaBNet. 24 |  25 | 26 | ## Dataset (Libri-SIM) 27 | The simulated dataset `Libri-SIM` is based on the [Librispeech](http://www.openslr.org/12) corpus. 28 | 29 | ### Raw data download 30 | 1) Download the *train-clean-100*, *dev-clean* and *test-clean* data from Librispeech's website and unzip them into any directory. The absolute path for the directory is denoted as *libri_path*, which should contain 3 subfolders *train-clean-100*, *dev-clean* and *test-clean*. 31 | 2) Download or clone this repository. 32 | 33 | ### Preprocess 34 | Run the command below for *train-clean-100*, *dev-clean* and *test-clean* respectively. 35 | ```python3 36 | python ./preparation/1_sb_vad.py 37 | python ./preparation/2_clip_segments.py 38 | python ./preparation/3_generate_list.py 39 | ``` 40 | The preprocessing here is to clip 4 seconds long speech segments with few silence for further simulation. 41 | 42 | 43 | ### Simulation 44 | - We simulate 6-channel audio data from original single-channel audio through pyroomacoustics, the spacings of 6 microphones are 0.04 m, 0.04 m, 0.12 m, 0.04 m, 0.04 m. The parameters of simulated rooms are shown in Tabel 1, the length of room is randomly selected between 4 m and 12 m, the width of room is a random number between 3 m and 9 m, and the height of room is a random number between 2.5 m and 5 m. There are small, middle, large 3 types of room according to the length of the room, the RT60 is a random number between 0.3 s and 0.6 s, 0.4 s and 0.7 s, 0.5 s and 0.8 s respectively. 45 | 46 |
47 | 
48 |
49 |
50 | - As shown in Figure 1, the microphone array is located in the middle of the wall, at a distance of 0.5 m from the wall and 2 m from the ground. In order to make sound sources cover the area in rooms better, we first set the direction-of-arrival of sound source, then we leave 0.5 m between the sound source and the microphone array and between the sound source and the wall, and divide the rest range into near, medium and far range, the distance between microphone array and sound source is a random number in 3 types of range, so we simulate one original single channel speech data at near, medium, far distance simultaneously.
51 | - We simulate 50 training rooms, 10 validation rooms and 10 testing rooms, respectively. The angular spacing between simulated sources is 1°.
52 | - Then we generate 3 sets of training, validation set, and testing set from the simulated rooms. We randomly select 2 sources by random angles and source-to-array distances.
53 | - The training set contains 40,000 utterances (44.44 hours), the validation set and the testing set contains 5,000 and 3,000 utterances, respectively.
54 |
55 | Run commands below to simulate and generate the `.lst` which contains dataset paths.
56 | ```python3
57 | python ./preparation/libri_sim_code/4_gen_simulated_room.py
58 | python ./preparation/libri_sim_code/5_generate_json.py
59 | python ./preparation/libri_sim_code/6_json_2_list.py
60 | ```
61 |
--------------------------------------------------------------------------------
/dataloader/dataloader.py:
--------------------------------------------------------------------------------
1 | import os
2 | import copy
3 | import json
4 | import librosa
5 | import numpy as np
6 | import multiprocessing as mp
7 | import soundfile as sf
8 | import torch
9 |
10 | import torch.utils.data as tud
11 | from torch.utils.data import Dataset
12 |
13 | def audioread(path, duration, fs = 16000):
14 | wave_data, sr = sf.read(path)
15 | if sr != fs:
16 | wave_data = librosa.resample(wave_data,sr,fs)
17 | return wave_data[:fs*duration]
18 |
19 |
20 | def activelev(data):
21 | eps = np.finfo(np.float32).eps
22 | max_val = (1. + eps) / (np.std(data)+eps)
23 | data = data * max_val
24 | return data
25 |
26 | def gaussian_func(gt_idx, output_dimention, sigma):
27 | indices = np.arange(output_dimention)
28 | out = np.array(np.exp(-1 * np.square(indices - gt_idx) / sigma ** 2))
29 | return out
30 |
31 | def encode_AS(input_list):
32 | # Azimuth Spectrum
33 | AS = []
34 | for input_idx in input_list:
35 | AS.append(gaussian_func(input_idx, 210, 8))
36 | AS = np.asarray(AS) # [1, 210]
37 | return AS
38 |
39 | def encode_DS(input_list):
40 | # Distance Spectrum
41 | DS = []
42 | for input_idx in input_list:
43 | DS.append(gaussian_func(input_idx, 200, 20))
44 | DS = np.asarray(DS)
45 | return DS
46 |
47 | def gaussian_func_2d(x_len=10, y_len=10, granularity=0.5, x_mu=5, y_mu=5, x_sigma_pow2=0.1, y_sigma_pow2=0.1):
48 | '''
49 | Input: x_mu, y_mu
50 | Output: 2D array (each element indicates the probability of active sound source at each coordinate)
51 | '''
52 | X, Y = np.meshgrid(np.arange(0, x_len, granularity), np.arange(0, y_len, granularity))
53 | print(f'X.shape {X.shape}')
54 | # TODO
55 | Z = np.array(np.exp(-0.5 * np.square(X - x_mu) / x_sigma_pow2 - 0.5 * np.square(Y - y_mu) / y_sigma_pow2))
56 | print(f'Z.shape {X.shape}')
57 |
58 | return Z.flatten() # [400]
59 |
60 | def encode_LS(x, y):
61 | # location heatmap
62 | DS = []
63 | DS.append(gaussian_func_2d(x_mu=x, y_mu=y))
64 | DS = np.asarray(DS) # [1, 400]
65 | return DS
66 |
67 | def parse_scp(scp, path_list):
68 | with open(scp, encoding='utf-8') as fid:
69 | for line in fid:
70 | path_list.append(line.strip())
71 |
72 | class TFDataset(Dataset):
73 | def __init__(self, wav_scp, data_mix_info, n_mics = 6, duration = 4, sample_rate= 16000,
74 | perturb_prob = 0.0, negatives=0.2, hold_coff=0.003, n_avb_mics = 2):
75 | mgr = mp.Manager()
76 | self.file_list = mgr.list()
77 | self.noise_list = mgr.list()
78 | self.index = mgr.list()
79 |
80 | self.data_mix_info = data_mix_info
81 | self.duration = duration
82 |
83 | self.sr = sample_rate
84 | self.n_mics = n_mics
85 | self.n_avb_mics = n_avb_mics
86 | self.perturb_prob = perturb_prob
87 | self.negatives = negatives
88 | self.hlod_coff = hold_coff
89 |
90 | self.angle_dimension = 210
91 | self.time_bins = 249
92 | self.speaker_num = 2
93 |
94 | pc_list = []
95 | p = mp.Process(target = parse_scp, args=(wav_scp,self.file_list))
96 | p.start()
97 | pc_list.append(p)
98 |
99 | for p in pc_list:
100 | p.join()
101 |
102 | self.index = [idx for idx in range(len(self.file_list))]
103 |
104 | def __len__(self):
105 | return len(self.index)
106 |
107 | def __getitem__(self, idx):
108 | item = self.index[idx]
109 | file_index = item
110 | file_path = self.file_list[file_index]
111 |
112 | with open(file_path) as json_file:
113 | metadata = json.load(json_file)
114 |
115 | all_sources, doa_as_array, target_1, target_2, wave_paths, doa_idx_array, xy_coordinates = self.get_mixture_and_gt(metadata)
116 |
117 | all_sources = torch.stack(all_sources,dim=0)
118 | mixed_data = torch.sum(all_sources,dim=0)
119 | channel_num, _ = mixed_data.size()
120 | target_1 = torch.stack(target_1,dim=0)
121 | target_1 = torch.sum(target_1,dim=0)
122 | target_2 = torch.stack(target_2,dim=0)
123 | target_2 = torch.sum(target_2,dim=0)
124 |
125 | scale = 0.5
126 | for channel_idx in range(channel_num):
127 | mix_single_channel_wav = mixed_data[channel_idx,:]
128 | max_amp = torch.max(torch.abs(mix_single_channel_wav))
129 | if max_amp == 0:
130 | max_amp =1
131 | mix_scale = 1/max_amp*scale
132 | mixed_data[channel_idx,:] = mixed_data[channel_idx,:] * mix_scale
133 | target_1[channel_idx, :] = target_1[channel_idx, :] * mix_scale
134 | target_2[channel_idx, :] = target_2[channel_idx, :] * mix_scale
135 |
136 | # without noise data
137 | target_1 = target_1[0, :]
138 | target_2 = target_2[0, :]
139 |
140 | return {
141 | "mixed_data": mixed_data,
142 | "target_1": target_1,
143 | "target_2": target_2,
144 | "doa_as_array": doa_as_array, # [B, T, S, n_mics, 210]
145 | "doa_idx_array": doa_idx_array, # [B, T, S, n_mics]
146 | "wave_paths": wave_paths,
147 | "xy": xy_coordinates # [B, T, S, 2]
148 | }
149 |
150 | # # with noise data (need test)
151 | # target_1 = target_1[0, :]
152 | # target_2 = target_1[0, :]
153 |
154 | # mixed_data = np.array(mixed_data)
155 | # target_1 = np.array(target_1)
156 | # target_2 = np.array(target_2)
157 | # interference = np.zeros(len(target_1))
158 |
159 | # stationary_noise_wav_path_part = metadata["stationary_noise"]['wave_path']
160 |
161 | # SNR = metadata["stationary_noise"]['SNR']
162 |
163 | # for channel_idx in range(channel_num):
164 | # stationary_noise_wav_path = stationary_noise_wav_path_part + '_' + str(channel_idx) + '.wav'
165 |
166 | # noise_change_weight = np.max(mixed_data)*(10 ** (-SNR / 20))
167 | # noise_wav = activelev(audioread(stationary_noise_wav_path, self.duration)) * noise_change_weight
168 | # mixed_data[channel_idx, :] = mixed_data[channel_idx, :] + noise_wav
169 |
170 | # if channel_idx == 0:
171 | # interference = interference + noise_wav
172 |
173 | # mixed_data = torch.from_numpy(mixed_data)
174 | # target_1 = torch.from_numpy(target_1)
175 | # target_2 = torch.from_numpy(target_2)
176 | # interference = torch.from_numpy(interference)
177 |
178 | # return mixed_data, doa_as_array, doa_ont_hot_array, target_1, target_2, interference
179 |
180 | def get_mixture_and_gt(self, metadata):
181 | # dataset_prefix = "/local02/fuyanjie"
182 | # dataset_prefix = "/sata/fuyanjie"
183 | all_sources = []
184 | source_index = 0
185 | target_data_1 = []
186 | target_data_2 = []
187 | wave_paths= {}
188 |
189 | as_dict = dict([])
190 | doa_idx_array = np.zeros([self.time_bins, self.speaker_num, self.n_avb_mics], dtype=np.int16)
191 | doa_as_array = np.zeros([self.time_bins, self.speaker_num, self.n_avb_mics, self.angle_dimension])
192 | angle_list = np.zeros([self.speaker_num], dtype=np.int16)
193 | xy_coordinates = np.zeros([self.time_bins, self.speaker_num, 2])
194 |
195 | for key in metadata.keys():
196 | if "source" in key:
197 | channel_index_list = np.arange(self.n_mics)
198 | flag = metadata[key]['wave_path']
199 | # flag.replace(dataset_prefix, "/CDShare3")
200 | gt_audio_files = [flag + '_'+ str(channel_index) + '.wav' for channel_index in channel_index_list]
201 | gt_waveforms = []
202 | for index, gt_audio_file in enumerate(gt_audio_files):
203 | gt_waveform = audioread(gt_audio_file, self.duration)
204 | single_channel_wav = activelev(gt_waveform)
205 | gt_waveforms.append(torch.from_numpy(single_channel_wav))
206 |
207 | shifted_gt = np.stack(gt_waveforms)
208 | perturbed_source = shifted_gt
209 |
210 | perturbed_source = torch.from_numpy(perturbed_source)
211 | perturbed_source = perturbed_source.to(torch.float32)
212 |
213 | if source_index !=0:
214 | # ignore the SIR between diffrent speakers
215 | # SIR = metadata[key]['SIR']
216 | SIR = 0
217 | change_weight = 10 ** (SIR/20)
218 | perturbed_source = perturbed_source * change_weight
219 |
220 | all_sources.append(perturbed_source)
221 |
222 | source_azimuth = int(metadata[key]['azimuth'])
223 | source2mic_dist = int(metadata[key]['s2m_dist'])
224 |
225 | source_azimuth1 = int(round(metadata[key]['azimuth1']))
226 | source_azimuth3 = int(round(metadata[key]['azimuth3']))
227 | source_azimuth4 = int(round(metadata[key]['azimuth4']))
228 | source_azimuth6 = int(round(metadata[key]['azimuth6']))
229 |
230 | if source_index == 0:
231 | target_data_1.append(perturbed_source)
232 | wave_paths["spk1"] = flag
233 | angle_list[source_index] = source_azimuth
234 | elif source_index == 1:
235 | target_data_2.append(perturbed_source)
236 | wave_paths["spk2"] = flag
237 | angle_list[source_index] = source_azimuth
238 |
239 | ASs = np.zeros([self.n_avb_mics, self.angle_dimension]) # [4, 210]
240 | if self.n_avb_mics == 2:
241 | ASs[0, :] = encode_AS([source_azimuth1 + 15])
242 | ASs[1, :] = encode_AS([source_azimuth6 + 15])
243 | elif self.n_avb_mics == 4:
244 | ASs[0, :] = encode_AS([source_azimuth1 + 15])
245 | ASs[1, :] = encode_AS([source_azimuth3 + 15])
246 | ASs[2, :] = encode_AS([source_azimuth4 + 15])
247 | ASs[3, :] = encode_AS([source_azimuth6 + 15])
248 | elif self.n_avb_mics == 1:
249 | ASs[0, :] = encode_AS([source_azimuth + 15])
250 |
251 | as_dict[source_index] = ASs
252 | vad_label = metadata[key]['vad_label']
253 | for vad_index in range(len(vad_label)-1):
254 | if vad_label[vad_index] == 1:
255 | source_azimuth_rad = source_azimuth / 180.0 * np.pi
256 | xy_coordinates[vad_index, source_index, 0] = source2mic_dist * np.cos(source_azimuth_rad)
257 | xy_coordinates[vad_index, source_index, 1] = source2mic_dist * np.sin(source_azimuth_rad)
258 | if self.n_avb_mics == 2:
259 | doa_idx_array[vad_index, source_index, 0] = source_azimuth1 + 15
260 | doa_idx_array[vad_index, source_index, 1] = source_azimuth6 + 15
261 | elif self.n_avb_mics == 4:
262 | doa_idx_array[vad_index, source_index, 0] = source_azimuth1 + 15
263 | doa_idx_array[vad_index, source_index, 1] = source_azimuth3 + 15
264 | doa_idx_array[vad_index, source_index, 2] = source_azimuth4 + 15
265 | doa_idx_array[vad_index, source_index, 3] = source_azimuth6 + 15
266 | elif self.n_avb_mics == 1:
267 | doa_idx_array[vad_index, source_index, 0] = source_azimuth + 15
268 | else:
269 | doa_idx_array[vad_index, source_index, :] = -1
270 | source_index = source_index + 1
271 |
272 |
273 | # Sort by Azimuth
274 | for i in range(0, self.speaker_num):
275 | for j in range(i+1, self.speaker_num):
276 | if angle_list[i] > angle_list[j]:
277 | temp_doa_arr_1 = copy.deepcopy(doa_idx_array[:, i, :])
278 | temp_doa_arr_2 = copy.deepcopy(doa_idx_array[:, j, :])
279 | doa_idx_array[:, i, :], doa_idx_array[:, j, :] = temp_doa_arr_2, temp_doa_arr_1
280 | temp_target_1 = copy.deepcopy(target_data_1)
281 | temp_target_2 = copy.deepcopy(target_data_2)
282 | target_data_1, target_data_2 = temp_target_2, temp_target_1
283 | temp1 = copy.deepcopy(angle_list[i])
284 | temp2 = copy.deepcopy(angle_list[j])
285 | angle_list[i], angle_list[j] = temp2, temp1
286 | temp_as_1 = copy.deepcopy(as_dict[i])
287 | temp_as_2 = copy.deepcopy(as_dict[j])
288 | as_dict[0], as_dict[1] = temp_as_2, temp_as_1
289 | temp_xy_1 = copy.deepcopy(xy_coordinates[:, i, :])
290 | temp_xy_2 = copy.deepcopy(xy_coordinates[:, j, :])
291 | xy_coordinates[:, i, :], xy_coordinates[:, j, :] = temp_xy_2, temp_xy_1
292 |
293 | # Assign after sorting
294 | for source_idx in range(0, source_index):
295 | for time_idx in range(0, self.time_bins):
296 | if (doa_idx_array[time_idx, source_idx, :] == -1).any():
297 | doa_as_array[time_idx, source_idx, :, :] = np.zeros([self.n_avb_mics, self.angle_dimension])
298 | else:
299 | azi_s = as_dict[source_idx]
300 | doa_as_array[time_idx, source_idx, :, :] = azi_s
301 |
302 | return all_sources, doa_as_array, target_data_1, target_data_2, wave_paths, doa_idx_array, xy_coordinates
303 |
304 | class Sampler(tud.sampler.Sampler):
305 | def __init__(self, data_source, batch_size):
306 | it_end = len(data_source) - batch_size + 1
307 | self.batches = [range(i,i+batch_size) for i in range(0, it_end, batch_size)]
308 | self.data_source = data_source
309 |
310 | def __iter__(self):
311 | return (i for b in self.batches for i in b)
312 |
313 | def __len__(self):
314 | return len(self.data_source)
315 |
316 | def static_loader(clean_scp, batch_size = 4, shuffle = True, num_workers = 8, duration = 4, sample_rate = 16000, data_mix_info = None, n_avb_mics = 2):
317 | dataset = TFDataset(
318 | wav_scp= clean_scp,
319 | data_mix_info = data_mix_info,
320 | duration = duration,
321 | sample_rate = sample_rate,
322 | n_avb_mics = n_avb_mics
323 | )
324 |
325 | loader = tud.DataLoader(
326 | dataset,
327 | batch_size = batch_size,
328 | shuffle = shuffle,
329 | num_workers = num_workers,
330 | drop_last = False
331 | )
332 | return loader
333 |
334 |
335 | if __name__ == "__main__":
336 | lst_path = '/Work21/2021/fuyanjie/pycode/LaBNet/data/exp_list/test-clean_0101.lst'
337 | data_loader = static_loader(lst_path, shuffle=True, batch_size=4)
338 | print(f'len(data_loader) {len(data_loader)}') # len(data_loader) is samples / batch_size
339 | one_batch_data = next(iter(data_loader))
340 | print('AAA ', one_batch_data["wave_paths"])
341 | print('B ', one_batch_data["doa_as_array"].dtype)
342 | print('C ', one_batch_data["doa_idx_array"].dtype)
343 | # AAA torch.Size([224, 12, 7, 257]) torch.Size([224])
344 | # one_batch_data = next(iter(data_loader))
345 | # print('BBB ', one_batch_data["clean_ris"].shape, one_batch_data["target_doa"].shape)
346 | os.makedirs('../debug_plot', exist_ok=True)
347 | import matplotlib.pyplot as plt
348 | for idx, data in enumerate(data_loader):
349 | fig, ax = plt.subplots() # 创建图实例
350 | ax.plot(np.linspace(0, 210, 210), data['doa_as_array'][3, 50, 0, 0], color='r')
351 | # ax.plot(np.linspace(0, 210, 210), data['doa_idx_array'][2, 50, 0, 0], color='b')
352 | print(f' AAA DOA {data["doa_as_array"].shape}')
353 | print(f' A doa_as_array {data["doa_as_array"][3, 50, 0, 0].shape}')
354 | print(f' B DOA {data["doa_idx_array"][3, 50, 0]}')
355 | plt.savefig(f'../debug_plot/test_as_spk1_{idx}.png')
356 | plt.cla()
357 | ax.plot(np.linspace(0, 210, 210), data['doa_as_array'][3, 50, 1, 0], color='r')
358 | # ax.plot(np.linspace(0, 210, 210), data['doa_one_hot_array'][2, 50, 1, 0], color='b')
359 | print(f' C doa_as_array {data["doa_as_array"][3, 50, 1, 0].shape}')
360 | print(f' D doa_idx_array {data["doa_idx_array"][3, 50, 1]}')
361 | plt.savefig(f'../debug_plot/test_as_spk2_{idx}.png')
362 | # plt.cla()
363 | # ax.plot(np.linspace(0, 400, 400), data['dist_ds_array'][2, 50, 1], color='r')
364 | # ax.plot(np.linspace(0, 400, 400), data['dist_one_hot_array'][2, 50, 1], color='b')
365 | # print(f' C Dist {data["dist_idx_array"][2, 50, 1]}')
366 | # print(f' D Dist {np.where(data["dist_one_hot_array"][2, 50, 1] == 1)[0]}')
367 | # plt.savefig(f'./test_ds_spk2_{idx}.png')
368 | # print(data['wave_paths'])
369 | # print(data['mixed_data'].shape)
370 | # print(data['target_1'].shape)
371 | # print(data['target_2'].shape)
372 | # print(data['doa_as_array'].shape)
373 | # print(data['doa_one_hot_array'].shape)
374 |
375 | print('------------')
376 | if idx >= 3:
377 | break
--------------------------------------------------------------------------------
/dataloader/test_dataloader.py:
--------------------------------------------------------------------------------
1 | import os
2 | import copy
3 | import json
4 | import librosa
5 | import numpy as np
6 | import multiprocessing as mp
7 | import soundfile as sf
8 | import torch
9 |
10 | import torch.utils.data as tud
11 | from torch.utils.data import Dataset
12 |
13 | def audioread(path, duration, fs = 16000):
14 | wave_data, sr = sf.read(path)
15 | if sr != fs:
16 | wave_data = librosa.resample(wave_data,sr,fs)
17 | return wave_data[:fs*duration]
18 |
19 |
20 | def activelev(data):
21 | eps = np.finfo(np.float32).eps
22 | max_val = (1. + eps) / (np.std(data)+eps)
23 | data = data * max_val
24 | return data
25 |
26 | def gaussian_func(gt_idx, output_dimention, sigma):
27 | indices = np.arange(output_dimention)
28 | out = np.array(np.exp(-1 * np.square(indices - gt_idx) / sigma ** 2))
29 | return out
30 |
31 | def encode_AS(input_list):
32 | # Azimuth Spectrum
33 | AS = []
34 | for input_idx in input_list:
35 | AS.append(gaussian_func(input_idx, 210, 8))
36 | AS = np.asarray(AS) # [1, 210]
37 | return AS
38 |
39 | def encode_DS(input_list):
40 | # Distance Spectrum
41 | DS = []
42 | for input_idx in input_list:
43 | DS.append(gaussian_func(input_idx, 200, 20))
44 | DS = np.asarray(DS)
45 | return DS
46 |
47 | def gaussian_func_2d(x_len=10, y_len=10, granularity=0.5, x_mu=5, y_mu=5, x_sigma_pow2=0.1, y_sigma_pow2=0.1):
48 | '''
49 | Input: x_mu, y_mu
50 | Output: 2D array (each element indicates the probability of active sound source at each coordinate)
51 | '''
52 | X, Y = np.meshgrid(np.arange(0, x_len, granularity), np.arange(0, y_len, granularity))
53 | print(f'X.shape {X.shape}')
54 | # TODO
55 | Z = np.array(np.exp(-0.5 * np.square(X - x_mu) / x_sigma_pow2 - 0.5 * np.square(Y - y_mu) / y_sigma_pow2))
56 | print(f'Z.shape {X.shape}')
57 |
58 | return Z.flatten() # [400]
59 |
60 | def encode_LS(x, y):
61 | # location heatmap
62 | DS = []
63 | DS.append(gaussian_func_2d(x_mu=x, y_mu=y))
64 | DS = np.asarray(DS) # [1, 400]
65 | return DS
66 |
67 | def parse_scp(scp, path_list):
68 | with open(scp, encoding='utf-8') as fid:
69 | # with open(scp, encoding='utf-8-sig') as fid:
70 | for line in fid:
71 | path_list.append(line.strip())
72 |
73 | class TFDataset(Dataset):
74 | def __init__(self, wav_scp, data_mix_info, n_mics = 6, duration = 4, sample_rate= 16000,
75 | perturb_prob = 0.0, negatives=0.2, hold_coff=0.003, n_avb_mics = 2):
76 |
77 | mgr = mp.Manager()
78 | self.file_list = mgr.list()
79 | self.noise_list = mgr.list()
80 | self.index = mgr.list()
81 |
82 | self.data_mix_info = data_mix_info
83 | self.duration = duration
84 |
85 | self.sr = sample_rate
86 | self.n_mics = n_mics
87 | self.n_avb_mics = n_avb_mics
88 | self.perturb_prob = perturb_prob
89 | self.negatives = negatives
90 | self.hlod_coff = hold_coff
91 |
92 | self.angle_dimension = 210
93 | self.time_bins = 249
94 | self.speaker_num = 2
95 |
96 | pc_list = []
97 | p = mp.Process(target = parse_scp, args=(wav_scp,self.file_list))
98 | p.start()
99 | pc_list.append(p)
100 |
101 | for p in pc_list:
102 | p.join()
103 |
104 | self.index = [idx for idx in range(len(self.file_list))]
105 |
106 | def __len__(self):
107 | return len(self.index)
108 |
109 | def __getitem__(self, idx):
110 | item = self.index[idx]
111 | file_index = item
112 | file_path = self.file_list[file_index]
113 |
114 | with open(file_path) as json_file:
115 | metadata = json.load(json_file)
116 |
117 | all_sources, doa_as_array, target_1, target_2, angle_list, wave_paths, doa_idx_array, xy_coordinates = self.get_mixture_and_gt(metadata)
118 |
119 | all_sources = torch.stack(all_sources,dim=0)
120 | mixed_data = torch.sum(all_sources,dim=0)
121 | channel_num, _ = mixed_data.size()
122 | target_1 = torch.stack(target_1,dim=0)
123 | target_1 = torch.sum(target_1,dim=0)
124 | target_2 = torch.stack(target_2,dim=0)
125 | target_2 = torch.sum(target_2,dim=0)
126 |
127 | scale = 0.5
128 | for channel_idx in range(channel_num):
129 | mix_single_channel_wav = mixed_data[channel_idx,:]
130 | max_amp = torch.max(torch.abs(mix_single_channel_wav))
131 | if max_amp == 0:
132 | max_amp =1
133 | mix_scale = 1/max_amp*scale
134 | mixed_data[channel_idx,:] = mixed_data[channel_idx,:] * mix_scale
135 | target_1[channel_idx, :] = target_1[channel_idx, :] * mix_scale
136 | target_2[channel_idx, :] = target_2[channel_idx, :] * mix_scale
137 |
138 | # without noise data
139 | target_1 = target_1[0, :]
140 | target_2 = target_2[0, :]
141 |
142 | angular_distance = abs(angle_list[0]-angle_list[1])
143 | social_distance = float(metadata["social_dist"])
144 |
145 | return {
146 | "mixed_data": mixed_data,
147 | "target_1": target_1,
148 | "target_2": target_2,
149 | "doa_as_array": doa_as_array, # [B, T, S, n_mics, 210]
150 | "doa_idx_array": doa_idx_array, # [B, T, S, n_mics]
151 | "angular_distance": angular_distance,
152 | "social_distance": social_distance,
153 | "wave_paths": wave_paths,
154 | "xy": xy_coordinates # [B, T, S, 2]
155 | }
156 |
157 | # # with noise data (need test)
158 | # target_1 = target_1[0, :]
159 | # target_2 = target_1[0, :]
160 |
161 | # mixed_data = np.array(mixed_data)
162 | # target_1 = np.array(target_1)
163 | # target_2 = np.array(target_2)
164 | # interference = np.zeros(len(target_1))
165 |
166 | # stationary_noise_wav_path_part = metadata["stationary_noise"]['wave_path']
167 |
168 | # SNR = metadata["stationary_noise"]['SNR']
169 |
170 | # for channel_idx in range(channel_num):
171 | # stationary_noise_wav_path = stationary_noise_wav_path_part + '_' + str(channel_idx) + '.wav'
172 |
173 | # noise_change_weight = np.max(mixed_data)*(10 ** (-SNR / 20))
174 | # noise_wav = activelev(audioread(stationary_noise_wav_path, self.duration)) * noise_change_weight
175 | # mixed_data[channel_idx, :] = mixed_data[channel_idx, :] + noise_wav
176 |
177 | # if channel_idx == 0:
178 | # interference = interference + noise_wav
179 |
180 | # mixed_data = torch.from_numpy(mixed_data)
181 | # target_1 = torch.from_numpy(target_1)
182 | # target_2 = torch.from_numpy(target_2)
183 | # interference = torch.from_numpy(interference)
184 |
185 | # return mixed_data, doa_as_array, doa_ont_hot_array, target_1, target_2, interference
186 |
187 | def get_mixture_and_gt(self, metadata):
188 | dataset_prefix = "/local01/fuyanjie"
189 | # dataset_prefix = "/sata/fuyanjie"
190 | all_sources = []
191 | source_index = 0
192 | target_data_1 = []
193 | target_data_2 = []
194 | wave_paths= {}
195 |
196 | as_dict = dict([])
197 | doa_idx_array = np.zeros([self.time_bins, self.speaker_num, self.n_avb_mics], dtype=np.int16)
198 | doa_as_array = np.zeros([self.time_bins, self.speaker_num, self.n_avb_mics, self.angle_dimension])
199 | angle_list = np.zeros([self.speaker_num], dtype=np.int16)
200 | xy_coordinates = np.zeros([self.time_bins, self.speaker_num, 2])
201 | for key in metadata.keys():
202 | if "source" in key:
203 | channel_index_list = np.arange(self.n_mics)
204 | flag = metadata[key]['wave_path']
205 | gt_audio_files = [flag + '_'+ str(channel_index) + '.wav' for channel_index in channel_index_list]
206 | gt_waveforms = []
207 | for index, gt_audio_file in enumerate(gt_audio_files):
208 | gt_waveform = audioread(gt_audio_file, self.duration)
209 | single_channel_wav = activelev(gt_waveform)
210 | gt_waveforms.append(torch.from_numpy(single_channel_wav))
211 |
212 | shifted_gt = np.stack(gt_waveforms)
213 | perturbed_source = shifted_gt
214 |
215 | perturbed_source = torch.from_numpy(perturbed_source)
216 | perturbed_source = perturbed_source.to(torch.float32)
217 |
218 | if source_index !=0:
219 | # ignore the SIR between diffrent speakers
220 | # SIR = metadata[key]['SIR']
221 | SIR = 0
222 | change_weight = 10 ** (SIR/20)
223 | perturbed_source = perturbed_source * change_weight
224 |
225 | all_sources.append(perturbed_source)
226 |
227 | source_azimuth = int(metadata[key]['azimuth'])
228 | source2mic_dist = int(metadata[key]['s2m_dist'])
229 |
230 | source_azimuth1 = int(round(metadata[key]['azimuth1']))
231 | source_azimuth3 = int(round(metadata[key]['azimuth3']))
232 | source_azimuth4 = int(round(metadata[key]['azimuth4']))
233 | source_azimuth6 = int(round(metadata[key]['azimuth6']))
234 |
235 | if source_index == 0:
236 | target_data_1.append(perturbed_source)
237 | wave_paths["spk1"] = flag
238 | angle_list[source_index] = source_azimuth
239 | elif source_index == 1:
240 | target_data_2.append(perturbed_source)
241 | wave_paths["spk2"] = flag
242 | angle_list[source_index] = source_azimuth
243 |
244 | ASs = np.zeros([self.n_avb_mics, self.angle_dimension]) # [4, 210]
245 | if self.n_avb_mics == 2:
246 | ASs[0, :] = encode_AS([source_azimuth1 + 15])
247 | ASs[1, :] = encode_AS([source_azimuth6 + 15])
248 | elif self.n_avb_mics == 4:
249 | ASs[0, :] = encode_AS([source_azimuth1 + 15])
250 | ASs[1, :] = encode_AS([source_azimuth3 + 15])
251 | ASs[2, :] = encode_AS([source_azimuth4 + 15])
252 | ASs[3, :] = encode_AS([source_azimuth6 + 15])
253 | elif self.n_avb_mics == 1:
254 | ASs[0, :] = encode_AS([source_azimuth + 15])
255 | as_dict[source_index] = ASs
256 |
257 | vad_label = metadata[key]['vad_label']
258 | for vad_index in range(len(vad_label)-1):
259 | if vad_label[vad_index] == 1:
260 | source_azimuth_rad = source_azimuth / 180.0 * np.pi
261 | xy_coordinates[vad_index, source_index, 0] = source2mic_dist * np.cos(source_azimuth_rad)
262 | xy_coordinates[vad_index, source_index, 1] = source2mic_dist * np.sin(source_azimuth_rad)
263 | if self.n_avb_mics == 2:
264 | doa_idx_array[vad_index, source_index, 0] = source_azimuth1 + 15
265 | doa_idx_array[vad_index, source_index, 1] = source_azimuth6 + 15
266 | elif self.n_avb_mics == 4:
267 | doa_idx_array[vad_index, source_index, 0] = source_azimuth1 + 15
268 | doa_idx_array[vad_index, source_index, 1] = source_azimuth3 + 15
269 | doa_idx_array[vad_index, source_index, 2] = source_azimuth4 + 15
270 | doa_idx_array[vad_index, source_index, 3] = source_azimuth6 + 15
271 | elif self.n_avb_mics == 1:
272 | doa_idx_array[vad_index, source_index, 0] = source_azimuth + 15
273 | else:
274 | doa_idx_array[vad_index, source_index, :] = -1
275 | source_index = source_index + 1
276 |
277 |
278 | # # Sort by Azimuth
279 | # for i in range(0, self.speaker_num):
280 | # for j in range(i+1, self.speaker_num):
281 | # if angle_list[i] > angle_list[j]:
282 | # temp_doa_arr_1 = copy.deepcopy(doa_idx_array[:, i, :])
283 | # temp_doa_arr_2 = copy.deepcopy(doa_idx_array[:, j, :])
284 | # doa_idx_array[:, i, :], doa_idx_array[:, j, :] = temp_doa_arr_2, temp_doa_arr_1
285 | # temp_target_1 = copy.deepcopy(target_data_1)
286 | # temp_target_2 = copy.deepcopy(target_data_2)
287 | # target_data_1, target_data_2 = temp_target_2, temp_target_1
288 | # temp1 = copy.deepcopy(angle_list[i])
289 | # temp2 = copy.deepcopy(angle_list[j])
290 | # angle_list[i], angle_list[j] = temp2, temp1
291 | # temp_as_1 = copy.deepcopy(as_dict[0])
292 | # temp_as_2 = copy.deepcopy(as_dict[1])
293 | # as_dict[0], as_dict[1] = temp_as_2, temp_as_1
294 |
295 | # Assign after sorting
296 | for source_idx in range(0, source_index):
297 | for time_idx in range(0, self.time_bins):
298 | if (doa_idx_array[time_idx, source_idx, :] == -1).any():
299 | doa_as_array[time_idx, source_idx, :, :] = np.zeros([self.n_avb_mics, self.angle_dimension])
300 | else:
301 | azi_s = as_dict[source_idx]
302 | doa_as_array[time_idx, source_idx, :, :] = azi_s
303 |
304 | return all_sources, doa_as_array, target_data_1, target_data_2, angle_list, wave_paths, doa_idx_array, xy_coordinates
305 |
306 | class Sampler(tud.sampler.Sampler):
307 | def __init__(self, data_source, batch_size):
308 | it_end = len(data_source) - batch_size + 1
309 | self.batches = [range(i,i+batch_size) for i in range(0, it_end, batch_size)]
310 | self.data_source = data_source
311 |
312 | def __iter__(self):
313 | return (i for b in self.batches for i in b)
314 |
315 | def __len__(self):
316 | return len(self.data_source)
317 |
318 | def static_loader(clean_scp, batch_size = 4, shuffle = True, num_workers = 8, duration = 4, sample_rate = 16000, data_mix_info = None, n_avb_mics = 2):
319 | dataset = TFDataset(
320 | wav_scp= clean_scp,
321 | data_mix_info = data_mix_info,
322 | duration = duration,
323 | sample_rate = sample_rate,
324 | n_avb_mics = n_avb_mics
325 | )
326 |
327 | loader = tud.DataLoader(
328 | dataset,
329 | batch_size = batch_size,
330 | shuffle = shuffle,
331 | num_workers = num_workers,
332 | drop_last = False
333 | )
334 | return loader
335 |
336 |
337 | if __name__ == "__main__":
338 | lst_path = '/Work21/2021/fuyanjie/pycode/LaBNet/data/exp_list/test-clean_0101.lst'
339 | data_loader = static_loader(lst_path, shuffle=False)
340 | print(f'len(data_loader) {len(data_loader)}') # len(data_loader) is samples / batch_size
341 | one_batch_data = next(iter(data_loader))
342 | print('AAA ', one_batch_data["wave_paths"])
343 | print('B ', one_batch_data["doa_as_array"].dtype)
344 | print('C ', one_batch_data["doa_idx_array"].dtype)
345 | # AAA torch.Size([224, 12, 7, 257]) torch.Size([224])
346 | # one_batch_data = next(iter(data_loader))
347 | # print('BBB ', one_batch_data["clean_ris"].shape, one_batch_data["target_doa"].shape)
348 | import matplotlib.pyplot as plt
349 | for idx, data in enumerate(data_loader):
350 | print(f' AAA DOA {data["doa_as_array"].shape}')
351 | print(f' A doa_as_array {data["doa_as_array"][3, 50, 0, 0].shape}')
352 | print(f' B DOA {data["doa_idx_array"][3, 50, 0]}')
353 |
354 | print(f' C doa_as_array {data["doa_as_array"][3, 50, 1, 0].shape}')
355 | print(f' D doa_idx_array {data["doa_idx_array"][3, 50, 1]}')
356 | # print(data['wave_paths'])
357 | # print(data['mixed_data'].shape)
358 | # print(data['target_1'].shape)
359 | # print(data['target_2'].shape)
360 | # print(data['doa_as_array'].shape)
361 | # print(data['doa_one_hot_array'].shape)
362 |
363 | print('------------')
364 | if idx >= 3:
365 | break
--------------------------------------------------------------------------------
/model/GRU.py:
--------------------------------------------------------------------------------
1 | import torch as th
2 | from torch.nn.utils.rnn import pad_packed_sequence, PackedSequence
3 |
4 | class RNN_MASK(th.nn.Module):
5 | def __init__(self,
6 | num_bins = 256, # num of feature channels
7 | freq_bins = 257,
8 | rnn = "gru",
9 | num_mask = 4,
10 | num_layer = 2,
11 | hidden_size = 500,
12 | dropout = 0.0,
13 | non_linear = "relu",
14 | bidirectional = False
15 | ):
16 | super(RNN_MASK,self).__init__()
17 | if non_linear not in ["relu","sigmoid","tanh"]:
18 | raise ValueError(
19 | "Unsupported non-linear type:{}".format(non_linear)
20 | )
21 | self.num_mask = num_mask
22 | rnn = rnn.upper()
23 | if rnn not in ["RNN", "LSTM", "GRU"]:
24 | raise ValueError(
25 | "Unsupported rnn type:{}".format(rnn)
26 | )
27 | self.rnn = getattr(th.nn, rnn)(
28 | num_bins,
29 | hidden_size,
30 | num_layer,
31 | batch_first = True,
32 | dropout = dropout,
33 | bidirectional = bidirectional
34 | )
35 | self.filter_size = 3
36 | self.drops = th.nn.Dropout(p=dropout)
37 | self.linear = th.nn.ModuleList([
38 | th.nn.Linear(hidden_size * 2 if bidirectional else hidden_size, 2 * freq_bins * self.filter_size * self.filter_size)
39 | for _ in range(self.num_mask)
40 | ])
41 |
42 | # self.conv1d = th.nn.Conv1d(hidden_size * 2 if bidirectional else hidden_size, self.num_mask * 2 * freq_bins * self.filter_size * self.filter_size, kernel_size=1)
43 |
44 | self.non_linear = {
45 | "relu": th.nn.functional.relu,
46 | "sigmoid": th.nn.functional.sigmoid,
47 | "tanh": th.nn.functional.tanh
48 | }[non_linear]
49 | self.num_bins = num_bins
50 |
51 | def forward(self,x):
52 | is_packed = isinstance(x, PackedSequence)
53 | if not is_packed and x.dim()!=3:
54 | x = th.unsqueeze(x,0)
55 | x, _ = self.rnn(x) # [B, C, T]
56 |
57 | if is_packed:
58 | x, _ = pad_packed_sequence(x,batch_first = True)
59 | x = self.drops(x)
60 | m = []
61 | for linear in self.linear:
62 | y = linear(x)
63 | y = self.non_linear(y)
64 | m.append(y)
65 | return m
--------------------------------------------------------------------------------
/model/Tree.py:
--------------------------------------------------------------------------------
1 | import sys
2 | sys.path.append("/Work21/2021/fuyanjie/pycode/LaBNetPro")
3 |
4 | from dataloader.dataloader import static_loader
5 | from model.GRU import RNN_MASK
6 | from torch_complex import ComplexTensor
7 | from torch_complex import functional as FC
8 | import torch.nn.functional as F
9 | import torch.nn as nn
10 | import torch
11 | import numpy as np
12 | from scipy.signal import get_window
13 |
14 | EPS = 1e-8
15 |
16 |
17 | def get_ipd(spec, dim):
18 | '''
19 | : param spec: (batch, channels, freq_bins:257 × 2, time_bins)
20 | : param dim:
21 | : return:
22 | '''
23 |
24 | real = spec[:, :, :dim, :]
25 | imag = spec[:, :, dim:, :]
26 | phase = torch.atan2(imag,real) # (batch,channels, freq_bins:257, time_bins)
27 | # 8, 12, 16, 16, 20 cm
28 | # diff1 = torch.cos((phase[:, 2, :, :] - phase[:, 0, :, :])).unsqueeze(1)
29 | # diff2 = torch.cos((phase[:, 3, :, :] - phase[:, 2, :, :])).unsqueeze(1)
30 | # diff3 = torch.cos((phase[:, 4, :, :] - phase[:, 2, :, :])).unsqueeze(1)
31 | # diff4 = torch.cos((phase[:, 3, :, :] - phase[:, 1, :, :])).unsqueeze(1)
32 | # diff5 = torch.cos((phase[:, 3, :, :] - phase[:, 0, :, :])).unsqueeze(1)
33 |
34 | # 8, 8, 12, 16, 16 cm
35 | # diff1 = torch.cos((phase[:, 2, :, :] - phase[:, 0, :, :])).unsqueeze(1)
36 | # diff2 = torch.cos((phase[:, 3, :, :] - phase[:, 1, :, :])).unsqueeze(1)
37 | # diff3 = torch.cos((phase[:, 3, :, :] - phase[:, 2, :, :])).unsqueeze(1)
38 | # diff4 = torch.cos((phase[:, 4, :, :] - phase[:, 2, :, :])).unsqueeze(1)
39 | # diff5 = torch.cos((phase[:, 5, :, :] - phase[:, 3, :, :])).unsqueeze(1)
40 |
41 | diff1 = torch.cos((phase[:, 1, :, :] - phase[:, 0, :, :])).unsqueeze(1)
42 | diff2 = torch.cos((phase[:, 2, :, :] - phase[:, 0, :, :])).unsqueeze(1)
43 | diff3 = torch.cos((phase[:, 3, :, :] - phase[:, 0, :, :])).unsqueeze(1)
44 | diff4 = torch.cos((phase[:, 4, :, :] - phase[:, 0, :, :])).unsqueeze(1)
45 | diff5 = torch.cos((phase[:, 5, :, :] - phase[:, 0, :, :])).unsqueeze(1)
46 |
47 | sin_diff1 = torch.sin((phase[:, 1, :, :] - phase[:, 0, :, :])).unsqueeze(1)
48 | sin_diff2 = torch.sin((phase[:, 2, :, :] - phase[:, 0, :, :])).unsqueeze(1)
49 | sin_diff3 = torch.sin((phase[:, 3, :, :] - phase[:, 0, :, :])).unsqueeze(1)
50 | sin_diff4 = torch.sin((phase[:, 4, :, :] - phase[:, 0, :, :])).unsqueeze(1)
51 | sin_diff5 = torch.sin((phase[:, 5, :, :] - phase[:, 0, :, :])).unsqueeze(1)
52 |
53 | result = torch.cat((diff1, diff2, diff3, diff4, diff5), dim=1)
54 | # result = torch.cat((diff1,diff2,diff3,diff4,diff5,sin_diff1,sin_diff2,sin_diff3,sin_diff4,sin_diff5), dim=1)
55 | #(batch, (m-1) pair microphones, freq bins:257, time_bins)
56 | return result
57 |
58 | def get_lps(spec, dim):
59 | '''
60 | : param spec: (batch,channels, freq _bins:257 × 2, time_bins)
61 | : param dim:
62 | : return:
63 | '''
64 | real = spec[:, :, :dim, :]
65 | imag = spec[:, :, dim:, :]
66 | mags = torch.sqrt(real ** 2 + imag ** 2)
67 | mags_refch = mags[:, 0, :, : ].unsqueeze(1)
68 | result = torch.log(mags_refch ** 2 + EPS) - np.log(EPS) # (batch, 1,freq_bins :257, time_bins)
69 | return result
70 |
71 | def get_spec_mag(spec, dim):
72 | '''
73 | :param spec: (batch, channels, freq bins:257 × 2, time_bins)
74 | :param dim:
75 | :return:
76 | '''
77 | real = spec[:, :, :dim, :]
78 | imag = spec[:, :, dim:, :]
79 | mags = torch.sqrt(real ** 2 + imag ** 2)
80 | mags_refch = mags[:, 0, :, :].unsqueeze(1)
81 | return mags_refch
82 |
83 |
84 | def get_covariance_v2(spec, crf, filter_size):
85 | batch, channels, two_freq_bins, time_bins = spec.size()
86 | freq_bins = two_freq_bins // 2
87 | spec = spec.reshape(batch, channels, 2, freq_bins, time_bins)
88 | spec = spec.permute(0, 1, 2, 4, 3)
89 | pad = nn.ZeroPad2d(1)
90 | spec_horizental = pad(spec)
91 | result_horizental = spec_horizental[:,:,:,:time_bins,:freq_bins+1]
92 | for i in range(1, filter_size):
93 | result_horizental = torch.cat((result_horizental, spec_horizental[:,:,:,i:i+time_bins,:freq_bins+1]),dim=3)
94 |
95 | spec_vertical = pad(result_horizental)
96 | result_vertical = spec_vertical[:,:,:,1:filter_size*time_bins+1,1:1+freq_bins]
97 | for j in range(1,filter_size):
98 | result_vertical = torch.cat((result_vertical,spec_vertical[:,:,:,1:3*time_bins+1,1+j:1+j+freq_bins]),dim=4)
99 |
100 | spec_final = result_vertical.reshape(batch, channels, 2, filter_size, time_bins, filter_size * freq_bins)
101 | spec_final = spec_final.reshape(batch, channels, 2, filter_size, time_bins, filter_size, freq_bins)
102 | spec_final = spec_final.permute(0, 1, 2, 4, 6, 3, 5)
103 |
104 | spec_real = spec_final[:,:,0,...]
105 | spec_imag = spec_final[:,:,1,...]
106 | spec_com = ComplexTensor(spec_real, spec_imag)
107 | spec_com = spec_com.permute(0, 3, 1, 2, 4, 5)
108 | crf_com = ComplexTensor(crf[:,0,...],crf[:,1,...])
109 | crf_com = crf_com.permute(0, 2, 1, 3, 4)
110 | psd_Y = FC.einsum("...ctkl,...etkl->...tcekl",[spec_com,spec_com.conj()])
111 | psd = psd_Y * crf_com[...,None,None,:,:]
112 | psd = psd.sum(dim=-1)
113 | psd = psd.sum(dim=-1)
114 |
115 | convariance_final_real = psd.real
116 | convariance_final_real = convariance_final_real.permute(0, 2, 1, 3, 4)
117 | convariance_final_imag = psd.imag
118 | convariance_final_imag = convariance_final_imag.permute(0, 2, 1, 3, 4)
119 |
120 | convariance_final = torch.stack((convariance_final_real,convariance_final_imag),dim=1)
121 |
122 | return convariance_final
123 |
124 | def remove_dc(signal):
125 | """Normalized to zero mean"""
126 | mean = torch.mean(signal, dim=-1, keepdim=True)
127 | signal = signal - mean
128 | return signal
129 |
130 | def dotproduct(y, y_hat) :
131 | #batch x channel x nsamples
132 | return torch.bmm(y.reshape(y.shape[0], 1, y.shape[ -1]), y_hat.reshape(y_hat.shape[0], y_hat.shape[-1], 1)).reshape(-1)
133 |
134 | def si_sdr_loss(e1, c1, c2):
135 | # [B, T]
136 | def sisdr(estimated, original, eps=1e-8):
137 | # estimated = remove_dc(estimated)
138 | # original = remove_dc(original)
139 | target = pow_norm(estimated, original) * original / (pow_p_norm(original) + eps)
140 | noise = estimated - target
141 | return -10 * torch.log10(eps + pow_p_norm(target) / pow_p_norm(noise) + eps)
142 |
143 | sisdr_loss = sisdr(e1, c1)
144 | avg_loss = torch.mean(sisdr_loss)
145 |
146 | return avg_loss
147 |
148 | def wsdr_loss(output, target_signal, inference_signal):
149 | """
150 | : param output: B, 1, T
151 | : param target_signal: (batch, time samples)
152 | : param inference_signal: (batch, time_samples)
153 | : return:
154 | """
155 |
156 | output = torch.squeeze(output, 1)
157 |
158 | y = target_signal
159 | z = inference_signal
160 |
161 | # target size: torch.Size([32, 15988B]) noise torch.Size([32, 159888]) output: torch.Size([32, 159800])# print ( "target size:", target_sig.size(),"noise ", noise.size(), 'output : ' , output .size(
162 | y_hat = output
163 | z_hat = y + z - y_hat # expected noise signal
164 |
165 | y_norm = torch.norm(y, dim=-1)
166 | z_norm = torch.norm(z, dim=-1)
167 | y_hat_norm = torch.norm(y_hat, dim=-1)
168 | z_hat_norm = torch.norm(z_hat, dim=-1)
169 |
170 | def loss_sdr(a, a_hat, a_norm, a_hat_norm):
171 | return dotproduct(a, a_hat) / (a_norm * a_hat_norm + EPS)
172 |
173 | alpha = y_norm.pow(2) / (y_norm.pow(2) + z_norm.pow(2) + EPS)
174 | loss_wSDR = -alpha * loss_sdr(y, y_hat, y_norm, y_hat_norm) - (1 - alpha) * loss_sdr(z, z_hat, z_norm, z_hat_norm)
175 |
176 | return loss_wSDR.mean()
177 |
178 | def pow_p_norm_np(signal):
179 | """Compute 2 Norm"""
180 | return np.square(np.linalg.norm(signal, ord=2, axis=-1))
181 |
182 | def pow_norm_np(s1, s2):
183 | return np.sum(s1 * s2, axis=-1)
184 |
185 | def pow_p_norm(signal):
186 | """Compute 2 Norm"""
187 | return torch.pow(torch.norm(signal, p=2, dim=-1, keepdim=True), 2)
188 |
189 | def pow_norm(s1, s2):
190 | return torch.sum(s1 * s2, dim=-1, keepdim=True)
191 |
192 | def pit_sisdr_loss(e1, e2, c1, c2):
193 | # [1, T]
194 | def sisdr(estimated, original):
195 | # estimated = remove_dc(estimated)
196 | # original = remove_dc(original)
197 | target = pow_norm(estimated, original) * original / pow_p_norm(original)
198 | noise = estimated - target
199 | return 10 * torch.log10(pow_p_norm(target) / pow_p_norm(noise))
200 |
201 | e1b = e1.squeeze() # [T]
202 | e2b = e2.squeeze() # [T]
203 | c1b = c1.squeeze() # [T]
204 | c2b = c2.squeeze() # [T]
205 |
206 | sdr1 = (sisdr(e1b, c1b) + sisdr(e2b, c2b)) * 0.5
207 | sdr2 = (sisdr(e2b, c1b) + sisdr(e1b, c2b)) * 0.5
208 |
209 | loss, idx = torch.max(torch.stack((sdr1, sdr2), dim=-1), dim=-1)
210 | avg_loss = torch.mean(loss)
211 |
212 | return avg_loss, idx
213 |
214 | def pit_sisdr_numpy(e1, e2, c1, c2):
215 | # [1, T]
216 | def sisdr(estimated, original):
217 | # estimated = remove_dc(estimated)
218 | # original = remove_dc(original)
219 | target = pow_norm_np(estimated, original) * original / pow_p_norm_np(original)
220 | noise = estimated - target
221 | return 10 * np.log10(pow_p_norm_np(target) / pow_p_norm_np(noise))
222 |
223 | e1b = np.squeeze(e1) # [T]
224 | e2b = np.squeeze(e2) # [T]
225 | c1b = np.squeeze(c1) # [T]
226 | c2b = np.squeeze(c2) # [T]
227 |
228 | sdr1 = (sisdr(e1b, c1b) + sisdr(e2b, c2b)) * 0.5
229 | sdr2 = (sisdr(e2b, c1b) + sisdr(e1b, c2b)) * 0.5
230 | print(f'sdr for permutation 1 {sdr1} sdr for permutation 2 {sdr2}', flush=True)
231 | loss = np.max(np.stack((sdr1, sdr2), axis=-1), axis=-1)
232 | idx = np.argmax(np.stack((sdr1, sdr2), axis=-1), axis=-1)
233 | avg_loss = np.mean(loss)
234 |
235 | return avg_loss, idx
236 |
237 | def init_kernels(win_len, win_inc, fft_len, win_type=None, invers=False):
238 | if win_type == 'None' or win_type is None:
239 | window = np.ones(win_len)
240 | else:
241 | window = get_window(win_type, win_len, fftbins=True) # **0.5
242 |
243 | N = fft_len
244 | fourier_basis = np.fft.rfft(np.eye(N))[:win_len]
245 | real_kernel = np.real(fourier_basis)
246 | imag_kernel = np.imag(fourier_basis)
247 | kernel = np.concatenate([real_kernel, imag_kernel], 1).T
248 |
249 | if invers:
250 | kernel = np.linalg.pinv(kernel).T
251 |
252 | kernel = kernel * window
253 | kernel = kernel[:, None, :]
254 | return torch.from_numpy(kernel.astype(np.float32)), torch.from_numpy(window[None, :, None].astype(np.float32))
255 |
256 |
257 | class ConvSTFT(nn.Module):
258 | def __init__(self, win_len, win_inc, fft_len=None, win_type='hamming', feature_type='real', fix=True):
259 | super(ConvSTFT, self).__init__()
260 |
261 | if fft_len == None:
262 | self.fft_len = np.int(2 ** np.ceil(np.log2(win_len)))
263 | else:
264 | self.fft_len = fft_len
265 |
266 | kernel, _ = init_kernels(win_len, win_inc, self.fft_len, win_type)
267 | # self.weight = nn.Parameter(kernel, requires_grad=(not fix))
268 | self.register_buffer('weight', kernel)
269 | self.feature_type = feature_type
270 | self.stride = win_inc
271 | self.win_len = win_len
272 | self.dim = self.fft_len
273 |
274 | def forward(self, inputs):
275 | if inputs.dim() == 2:
276 | inputs = torch.unsqueeze(inputs, 1)
277 | inputs = F.pad(inputs, [self.win_len - self.stride, self.win_len - self.stride])
278 | outputs = F.conv1d(inputs, self.weight, stride=self.stride)
279 |
280 | if self.feature_type == 'complex':
281 | return outputs
282 | else:
283 | dim = self.dim // 2 + 1
284 | real = outputs[:, :dim, :]
285 | imag = outputs[:, dim:, :]
286 | mags = torch.sqrt(real ** 2 + imag ** 2)
287 | phase = torch.atan2(imag, real)
288 | return mags, phase
289 |
290 |
291 | class ConviSTFT(nn.Module):
292 | def __init__(self, win_len, win_inc, fft_len=None, win_type='hamming', feature_type='real', fix=True):
293 | super(ConviSTFT, self).__init__()
294 | if fft_len == None:
295 | self.fft_len = np.int(2 ** np.ceil(np.log2(win_len)))
296 | else:
297 | self.fft_len = fft_len
298 | kernel, window = init_kernels(win_len, win_inc, self.fft_len, win_type, invers=True)
299 | # self.weight = nn.Parameter(kernel, requires_grad=(not fix))
300 | self.register_buffer('weight', kernel)
301 | self.feature_type = feature_type
302 | self.win_type = win_type
303 | self.win_len = win_len
304 | self.stride = win_inc
305 | self.stride = win_inc
306 | self.dim = self.fft_len
307 | self.register_buffer('window', window)
308 | self.register_buffer('enframe', torch.eye(win_len)[:, None, :])
309 |
310 | def forward(self, inputs, phase=None):
311 | """
312 | inputs : [B, N+2, T] (complex spec) or [B, N//2+1, T] (mags)
313 | phase: [B, N//2+1, T] (if not none)
314 | """
315 |
316 | if phase is not None:
317 | real = inputs * torch.cos(phase)
318 | imag = inputs * torch.sin(phase)
319 | inputs = torch.cat([real, imag], 1)
320 | outputs = F.conv_transpose1d(inputs, self.weight, stride=self.stride)
321 |
322 | # this is from torch-stft: https://github.com/pseeth/torch-stft
323 | t = self.window.repeat(1, 1, inputs.size(-1)) ** 2
324 | coff = F.conv_transpose1d(t, self.enframe, stride=self.stride)
325 | outputs = outputs / (coff + 1e-8)
326 | # outputs = torch.where(coff == 0, outputs, outputs/coff)
327 | outputs = outputs[..., self.win_len - self.stride:-(self.win_len - self.stride)]
328 |
329 | return outputs
330 |
331 | class Triangulation(nn.Module):
332 | def __init__(self):
333 | super(Triangulation, self).__init__()
334 |
335 | def forward(self, doas, mic_itval=28):
336 | """
337 | doas : [B*F, T, n_avb_mics]
338 | xy: [B*F, T, 2]
339 | """
340 | if doas.shape[-1] == 4:
341 | doas = torch.index_select(doas, -1, torch.tensor([0, 3], device=doas.device)) # [B*F, T, 2]
342 |
343 | # doas = torch.max(azi_s, dim=-1)[1] # [B, T, n_avb_mics] if input [B, T, n_avb_mics, 210]
344 | # if azi_s.shape[-1] == 4:
345 | # doas = torch.index_select(azi_s, -1, torch.tensor([0, 3])) # [B, T, 2]
346 |
347 | doa1 = doas[:, :, 0]
348 | doa2 = doas[:, :, 1]
349 | if (doa1 == doa2).any():
350 | doa1 = doa1 + 1
351 | doa1 = doa1 / 180.0 * np.pi
352 | doa2 = doa2 / 180.0 * np.pi
353 |
354 | side = torch.clamp(torch.sin(doa2) * mic_itval / (torch.sin(doa1-doa2) + 1e-8), min=50, max=800)
355 |
356 | x = -side * torch.cos(doa1)
357 | y = side * torch.sin(doa1)
358 | xy = torch.stack((x, y), dim=2)
359 | # print(f'side {side} doa1 {doa1} doa2 {doa2} xy {xy}', flush=True)
360 | return xy
361 |
362 |
363 | class Tree(nn.Module):
364 | def __init__(self, fft_len=512,speaker_num=2,n_avb_mics=2):
365 | super(Tree, self).__init__()
366 | self.channels = 6
367 | self.dim = fft_len // 2 + 1
368 | self.win_len = fft_len
369 | self.win_inc = fft_len // 2
370 | self.fft_len = fft_len
371 | self.win_type = 'hamming'
372 | self.filter_size = 3
373 | self.stftConv = ConvSTFT(self.win_len, self.win_inc, self.fft_len, self.win_type, feature_type='complex')
374 | self.istftConv = ConviSTFT(self.win_len, self.win_inc, self.fft_len, self.win_type, feature_type='complex')
375 | self.P = 3
376 | self.R = 4
377 | self.X = 8
378 | self.causal = True
379 | self.norm_type = 'cLN'
380 | self.n_mics = 6 # len(mic_location)
381 | self.n_avb_mics = n_avb_mics
382 | self.speaker_num = speaker_num
383 | self.reduction_linear = nn.Linear(in_features = (10+1)*257,out_features=256)
384 |
385 | self.repeats = RNN_MASK()
386 |
387 | self.width_window_length = 2
388 |
389 | self.RELU = nn.ReLU()
390 | self.PRELU = nn.PReLU()
391 | self.Sigmoid = nn.Sigmoid()
392 | self.covariance_ln_ss = nn.LayerNorm([self.n_mics, self.n_mics])
393 | self.covariance_ln_nn = nn.LayerNorm([self.n_mics, self.n_mics])
394 |
395 | self.conv1_doa = nn.Sequential(
396 | nn.Conv2d(144, 210*self.n_avb_mics, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0)),
397 | nn.BatchNorm2d(210*self.n_avb_mics), nn.ReLU(inplace=True)
398 | )
399 | self.conv_emb_doa = nn.Sequential(
400 | nn.Conv2d(210, 1, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0)),
401 | nn.BatchNorm2d(1), nn.ReLU(inplace=True)
402 | )
403 | self.conv2_doa = nn.Sequential(
404 | nn.Conv2d(257, 1, kernel_size=(3, 5), stride=(1, 1), padding=(1, 2)),
405 | nn.BatchNorm2d(1), nn.ReLU(inplace=True)
406 | )
407 |
408 | self.gru_doa = nn.GRU(210, hidden_size=210, num_layers=2, batch_first=True)
409 |
410 | ###
411 | # self.triangulation = nn.Linear(self.n_avb_mics * 210, 400)
412 | self.triangulation = Triangulation()
413 | # self.conv1_loc = nn.Sequential(
414 | # nn.Conv2d(210*self.n_avb_mics, 400, kernel_size=(3, 5), stride=(1, 1), padding=(1, 2)),
415 | # nn.BatchNorm2d(400), nn.ReLU(inplace=True)
416 | # )
417 | # self.conv2_loc = nn.Sequential(
418 | # nn.Conv2d(257, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
419 | # nn.BatchNorm2d(1), nn.ReLU(inplace=True)
420 | # )
421 | ###
422 |
423 | # Branch 1 (gru_bss)
424 | self.linear_bss_1 = nn.Linear(self.n_mics * self.n_mics * 2 * 2 + self.n_avb_mics * 210 + 2, 300)
425 | # self.linear_bss_1 = nn.Linear(self.n_mics * self.n_mics * 2 * 2 + self.n_avb_mics * 210, 300)
426 | self.gru_bss_1 = nn.GRU(300, hidden_size=300, num_layers=2, batch_first=True)
427 | self.linear_w1 = nn.Linear(300, self.channels * 2)
428 |
429 | # Branch 2 (gru_bss)
430 | self.linear_bss_2 = nn.Linear(self.n_mics * self.n_mics * 2 * 2 + self.n_avb_mics * 210 + 2, 300)
431 | # self.linear_bss_2 = nn.Linear(self.n_mics * self.n_mics * 2 * 2 + self.n_avb_mics * 210, 300)
432 | self.gru_bss_2 = nn.GRU(300, hidden_size=300, num_layers=2, batch_first=True)
433 | self.linear_w2 = nn.Linear(300, self.channels * 2)
434 |
435 | def get_params(self, weight_decay=0.0):
436 | #add L2 penalty
437 | weights, biases = [], []
438 | for name, param in self.named_parameters():
439 | if 'bias' in name:
440 | biases += [param]
441 | else:
442 | weights += [param]
443 | params = [{
444 | 'params': weights,
445 | 'weight_decay': weight_decay,
446 | }, {
447 | 'params': biases,
448 | 'weight_decay': 0.0,
449 | }]
450 | return params
451 |
452 | def forward(self, inputs):
453 | """
454 | Returns
455 | -------
456 | azis_1 / azis_2: [B, T, n_mics=self.n_avb_mics, 210]
457 | es_sig: [B, 1, T]
458 | """
459 | self.gru_doa.flatten_parameters()
460 | self.gru_bss_1.flatten_parameters()
461 | self.gru_bss_2.flatten_parameters()
462 | inputs = torch.unsqueeze(inputs, 2) # bs * channels * 1 * time_samples
463 | batch, channels, _, time_samples = inputs.size()
464 | inputs = inputs.view(batch * channels, 1, time_samples) # bs * channels, l, time_samples
465 |
466 | #STFT
467 | #que stftConv返回的维度是2 * 257吗
468 | spectrograms = self.stftConv(inputs)
469 | _, double_freq_bins, time_bins = spectrograms.size() # batch * channels, 2 * 257,time bins
470 | freq_bins = double_freq_bins // 2
471 | spectrograms = spectrograms.reshape(batch, channels, double_freq_bins, time_bins)
472 |
473 | '''compute ipd , lps and width feature,concat'''
474 | ipd = get_ipd(spectrograms, self.dim) # batch, M microphone pair num, freq_bins, time_bin$
475 | lps = get_spec_mag(spectrograms, self.dim) # batch, 1, freq_bins, time_bins
476 |
477 | audio_blocks = torch.cat((ipd, lps), dim=1) #(batch, n_mics, freq_bins:257, time_bins
478 | audio_blocks = audio_blocks.view(batch, (10+1) * 257, time_bins)
479 | audio_blocks = audio_blocks.permute(0,2,1) # [B, time_bins, 6*freq_bins]
480 |
481 | '''pass through front-GRU '''
482 | audio_reduction = self.reduction_linear(audio_blocks) # batch, time_bins, 256 channels
483 | masks = self.repeats(audio_reduction) # batch, time_bins, 256 channels
484 |
485 | convariance_set_doa_list = []
486 | convariance_set_bss_list = []
487 |
488 | """audio embeddings经过blstm得到 crf,crf与spectrograms得到目标信号的协方差矩阵"""
489 |
490 | for mask_idx in range(self.speaker_num):
491 | comlex_filter_ss = masks[mask_idx * 2].permute(0,2,1)
492 | comlex_filter_nn = masks[mask_idx * 2 + 1].permute(0,2,1)
493 |
494 | comlex_filter_ss = comlex_filter_ss.view(batch, freq_bins, 2 * self.filter_size * self.filter_size, time_bins)
495 | comlex_filter_ss = comlex_filter_ss.view(batch, freq_bins, 2, self.filter_size * self.filter_size, time_bins)
496 | comlex_filter_ss = comlex_filter_ss.view(batch, freq_bins, 2, self.filter_size, self.filter_size, time_bins)
497 | comlex_filter_ss = comlex_filter_ss.permute(0, 2, 5, 1, 3, 4)
498 | covariance_ss = get_covariance_v2(spectrograms, comlex_filter_ss, self.filter_size)
499 | covariance_ss = self.covariance_ln_ss(covariance_ss)
500 | covariance_ss = covariance_ss.unsqueeze(dim=4)
501 |
502 | comlex_filter_nn = comlex_filter_nn.view(batch, freq_bins, 2 * self.filter_size * self.filter_size, time_bins)
503 | comlex_filter_nn = comlex_filter_nn.view(batch, freq_bins, 2, self.filter_size * self.filter_size, time_bins)
504 | comlex_filter_nn = comlex_filter_nn.view(batch, freq_bins, 2, self.filter_size, self.filter_size, time_bins)
505 | comlex_filter_nn = comlex_filter_nn.permute(0, 2, 5, 1, 3, 4)
506 | covariance_nn = get_covariance_v2(spectrograms, comlex_filter_nn, self.filter_size)
507 | covariance_nn = self.covariance_ln_nn(covariance_nn)
508 | covariance_nn = covariance_nn.unsqueeze(dim=4)
509 |
510 | # TODO
511 | covariance_set = torch.cat((covariance_ss, covariance_nn), dim=4) # B,2,T,F,2,6,6
512 |
513 | covariance_set = covariance_set.permute(2, 0, 3, 1, 4, 5, 6)
514 | # (time bins, batch, freq bins, 2, 2, channels, channels)
515 | covariance_set = covariance_set.reshape(time_bins, batch, freq_bins, 2, 2, channels *channels)
516 | covariance_set = covariance_set.reshape(time_bins, batch, freq_bins, 2, 2* channels *channels)
517 | covariance_set = covariance_set.reshape(time_bins, batch, freq_bins, 2* 2* channels *channels)
518 |
519 | covariance_set_doa = covariance_set
520 | covariance_set_doa = covariance_set_doa.permute(1,3,0,2)
521 | convariance_set_doa_list.append(covariance_set_doa) # (B, 144, T, F)
522 |
523 | covariance_set = covariance_set.reshape(time_bins, batch * freq_bins, 2* 2* channels *channels)
524 | covariance_set = covariance_set.permute(1,0,2) # [B*F, T, 144]
525 |
526 | convariance_set_bss_list.append(covariance_set) # [B*F, T, 144]
527 |
528 | '''Branch_1'''
529 | # DoA_1
530 | # as1_freq = self.conv1_doa(convariance_set_doa_list[0]) # [B, 210*n_avb_mics, T, F]
531 | # as1_loc_emb = as1_freq.permute(0, 3, 2, 1) # [B, F, T, 210*n_avb_mics]
532 | # as1_freq = as1_freq.reshape(batch, 210, self.n_avb_mics, time_bins, freq_bins) # [B, 210, n_avb_mics, T, F]
533 | # as1_freq = as1_freq.reshape(batch, 210, self.n_avb_mics*time_bins, freq_bins) # [B, 210, T*n_avb_mics, F]
534 | # as1_freq_emb = as1_freq # [B, 210, T*n_avb_mics, F]
535 | # as1_freq = as1_freq.reshape(batch, 210, time_bins, self.n_avb_mics, freq_bins)
536 | # as1_freq = as1_freq.reshape(batch*self.n_avb_mics, 210, time_bins, freq_bins)
537 | # as1_freq = as1_freq.permute(0, 3, 2, 1) # [B*n_avb_mics, F, T, 210]
538 | # as1_freq_emb = self.conv_emb_doa(as1_freq_emb) # [B, 1, T*n_avb_mics, F]
539 | # as1_freq_emb = as1_freq_emb.squeeze(dim=1) # [B, T*n_avb_mics, F]
540 | # as1_freq_emb = as1_freq_emb.reshape(batch, time_bins, freq_bins, self.n_avb_mics) # [B, T, F, n_avb_mics]
541 | # as1_freq_emb = as1_freq_emb.reshape(batch * freq_bins, time_bins, self.n_avb_mics) # [B*F, T, n_avb_mics]
542 |
543 | ###
544 | as1_freq = self.conv1_doa(convariance_set_doa_list[0]) # [B, 210*n_avb_mics, T, F]
545 | as1_freq = as1_freq.permute(0, 3, 2, 1) # [B, F, T, 210*n_avb_mics]
546 | as1_freq_emb = as1_freq.reshape(batch * freq_bins, time_bins, 210*self.n_avb_mics) # [B*F, T, 210*n_avb_mics]
547 | as1_freq = as1_freq.reshape(batch, freq_bins, time_bins, 210, self.n_avb_mics)
548 | as1_freq = as1_freq.reshape(batch*self.n_avb_mics, freq_bins, time_bins, 210) # [B*n_avb_mics, F, T, 210]
549 |
550 | as_1 = self.conv2_doa(as1_freq) # (B*n_avb_mics, 1, T, 210)
551 | as_1 = as_1.squeeze(dim=1) # [B*n_avb_mics, T, 210]
552 | azis_1, _ = self.gru_doa(as_1) # [B*n_avb_mics, T, 210]
553 | azis_1 = azis_1.reshape(batch, self.n_avb_mics, time_bins, 210) # [B, n_avb_mics, T, 210]
554 | azis_1 = azis_1.permute(0, 2, 1, 3) # [B, T, n_mics=n_avb_mics, 210]
555 | doas_1 = torch.max(azis_1, dim=3)[1]
556 | xy1 = self.triangulation(doas_1) # [B, T, 2]
557 | xy1_emb = torch.repeat_interleave(xy1.unsqueeze(dim=1),repeats=257,dim=1) # [B, F, T, 2]
558 | xy1_emb = xy1_emb.reshape(batch*freq_bins, time_bins, 2) # [B*F, T, 2]
559 |
560 | # BSS_1
561 | convariance_set_with_sps_1 = torch.cat((convariance_set_bss_list[0], as1_freq_emb, xy1_emb), dim=2)
562 | # convariance_set_with_sps_1 = torch.cat((convariance_set_bss_list[0], as1_freq_emb), dim=2)
563 | covariance_set_bss_1 = self.linear_bss_1(convariance_set_with_sps_1) # [B*F, T, 300]
564 | gru_output_w1, _ = self.gru_bss_1(covariance_set_bss_1) # [B*F, T, 300]
565 | gru_output_w1 = gru_output_w1.reshape(batch, freq_bins, time_bins, 300)
566 | gru_output_w1 = gru_output_w1.permute(2, 0, 1, 3)
567 | w1 = self.linear_w1(gru_output_w1) # (time_bins, batch, freq_bins, channels * 2)
568 | w1 = w1.reshape(time_bins, batch, freq_bins, channels, 2) # (time_bins, batch, freq_bins, channels, 2)
569 | beamformer_1 = w1.permute(1, 0, 4, 2, 3) # (batch, time_bins, 2, freq_bins, channels)
570 | beamformer_1 = beamformer_1.unsqueeze(dim=5) # (batch, time_bins, 2, freq_bins, channels, 1)
571 |
572 | '''Branch_2'''
573 | # DoA_2
574 | # as2_freq = self.conv1_doa(convariance_set_doa_list[1]) # [B, 210*n_avb_mics, T, F]
575 | # as2_freq = as2_freq.reshape(batch, 210, self.n_avb_mics, time_bins, freq_bins)
576 | # as2_freq = as2_freq.reshape(batch, 210, self.n_avb_mics*time_bins, freq_bins) # [B, 210, T*n_avb_mics, F]
577 | # as2_freq_emb = as2_freq # [B, 210, T*n_avb_mics, F]
578 | # as2_freq = as2_freq.reshape(batch, 210, time_bins, self.n_avb_mics, freq_bins)
579 | # as2_freq = as2_freq.reshape(batch*self.n_avb_mics, 210, time_bins, freq_bins)
580 | # as2_freq = as2_freq.permute(0, 3, 2, 1) # [B*n_avb_mics, F, T, 210]
581 | # as2_freq_emb = self.conv_emb_doa(as2_freq_emb) # [B, 1, T*n_avb_mics, F]
582 | # as2_freq_emb = as2_freq_emb.squeeze(dim=1) # [B, T*n_avb_mics, F]
583 | # as2_freq_emb = as2_freq_emb.reshape(batch, time_bins, freq_bins, self.n_avb_mics) # [B, T, F, n_avb_mics]
584 | # as2_freq_emb = as2_freq_emb.reshape(batch * freq_bins, time_bins, self.n_avb_mics) # [B*F, T, n_avb_mics]
585 |
586 | ###
587 | as2_freq = self.conv1_doa(convariance_set_doa_list[1]) # [B, 210*n_avb_mics, T, F]
588 | as2_freq = as2_freq.permute(0, 3, 2, 1) # [B, F, T, 210*n_avb_mics]
589 | as2_freq_emb = as2_freq.reshape(batch * freq_bins, time_bins, 210*self.n_avb_mics) # [B*F, T, 210*n_avb_mics]
590 | as2_freq = as2_freq.reshape(batch, freq_bins, time_bins, 210, self.n_avb_mics)
591 | as2_freq = as2_freq.reshape(batch*self.n_avb_mics, freq_bins, time_bins, 210) # [B*n_avb_mics, F, T, 210]
592 |
593 | as_2 = self.conv2_doa(as2_freq) # (B*n_avb_mics, 1, T, 210)
594 | as_2 = as_2.squeeze(dim=1) # [B*n_avb_mics, T, 210]
595 | azis_2, _ = self.gru_doa(as_2) # [B*n_avb_mics, T, 210]
596 | azis_2 = azis_2.reshape(batch, self.n_avb_mics, time_bins, 210) # [B, n_avb_mics, T, 210]
597 | azis_2 = azis_2.permute(0, 2, 1, 3) # [B, T, n_mics=n_avb_mics, 210]
598 | doas_2 = torch.max(azis_2, dim=3)[1]
599 | xy2 = self.triangulation(doas_2) # [B, T, 2]
600 | xy2_emb = torch.repeat_interleave(xy2.unsqueeze(dim=1),repeats=257,dim=1) # [B, F, T, 2]
601 | xy2_emb = xy2_emb.reshape(batch*freq_bins, time_bins, 2) # [B*F, T, 2]
602 |
603 | # BSS_2
604 | convariance_set_with_sps_2 = torch.cat((convariance_set_bss_list[1], as2_freq_emb, xy2_emb), dim=2)
605 | # convariance_set_with_sps_2 = torch.cat((convariance_set_bss_list[1], as2_freq_emb), dim=2)
606 | covariance_set_bss_2 = self.linear_bss_2(convariance_set_with_sps_2)
607 | gru_output_w2, _ = self.gru_bss_2(covariance_set_bss_2)
608 | gru_output_w2 = gru_output_w2.reshape(batch, freq_bins, time_bins, 300)
609 | gru_output_w2 = gru_output_w2.permute(2, 0, 1, 3)
610 | w2 = self.linear_w2(gru_output_w2) # (time_bins, batch, freq_bins, channels * 2)
611 | w2 = w2.reshape(time_bins, batch, freq_bins, channels, 2) # (time_bins, batch, freq_bins, channels, 2)
612 | beamformer_2 = w2.permute(1, 0, 4, 2, 3) # (batch, time_bins, 2, freq_bins, channels)
613 | beamformer_2 = beamformer_2.unsqueeze(dim=5) # (batch, time_bins, 2, freq_bins,channels, 1)
614 |
615 | '''Beamformer Separation'''
616 | spec = torch.unsqueeze(spectrograms, 4) # (batch, channels, 2*freq_bins,time_bins, 1)
617 | # (batch, channels, 2, freq_bins, time bins, 1)
618 | spec = spec.reshape(batch, channels, 2, freq_bins, time_bins, 1)
619 | spec = spec.permute(0, 4, 2, 3, 1, 5) # (batch, time bins, 2, freq_bins, channels, 1)
620 | spec_real = spec[:, :, 0, :, :, :]# (batch, time bins, freq bins, channels, 1)
621 | spec_imag = spec[:, :, 1, :, :, :] # (batch,time bins, freq bins,channels, 1)
622 |
623 | beamformer_list = []
624 | extract_signal_list = []
625 | beamformer_list.append(beamformer_1)
626 | beamformer_list.append(beamformer_2)
627 |
628 | for beamformer in beamformer_list:
629 | beamformer_real = beamformer[:, :, 0, ...] # (batch,time bins, freq_bins,channels, 1)
630 | beamformer_imag = beamformer[:, :, 1, ...] # (batch, time bins, freq bins,channels, 1)
631 | beamformer_real = beamformer_real.permute(0, 1, 2, 4, 3) # (batch, time_ bins, freq_bins,1, channels)
632 | beamformer_imag = beamformer_imag.permute(0, 1, 2, 4, 3) #(batch, time_bins, freqg_bins, 1, channels)
633 |
634 | enhancement_real = torch.matmul(beamformer_real, spec_real) - torch.matmul(beamformer_imag, spec_imag)
635 | enhancement_imag = torch.matmul(beamformer_real, spec_imag) + torch.matmul(beamformer_imag, spec_real)#( batch, time bins, freq bins, 1, 1)
636 |
637 | enhancement_real = enhancement_real.squeeze(dim=-1) # (batch, time_bins, freq_bins, 1)
638 | enhancement_imag = enhancement_imag.squeeze(dim=-1) # (batch, time_bins,freq_bins, 1)
639 |
640 | enhancement_real = enhancement_real.permute(0, 1, 3, 2)
641 | enhancement_imag = enhancement_imag.permute(0, 1, 3, 2)
642 |
643 | enhancement = torch.cat((enhancement_real, enhancement_imag), dim=2) # batch, time_bins, 2, freq_bins
644 | enhancement = enhancement. reshape(batch, time_bins, 2 * freq_bins) # batch, time bins, 2 * freq bins
645 | enhancement = enhancement.permute(0, 2, 1) # batch, 2 * freq bins, time_bins
646 | '''do iSTFT'''
647 | extract_signal = self.istftConv(enhancement) # batch, 1, time_samples
648 | extract_signal_list.append(extract_signal)
649 |
650 | extract_signal_1 = extract_signal_list[0] # [B, 1, T]
651 | extract_signal_2 = extract_signal_list[1]
652 |
653 | return azis_1, azis_2, extract_signal_1, extract_signal_2
654 | # return azis_1, azis_2, xy1, xy2, extract_signal_1, extract_signal_2
655 |
656 | if __name__ == '__main__':
657 | # e1 = np.random.randn(4, 32000)
658 | # e2 = np.random.randn(4, 32000)
659 | # c1 = np.random.randn(4, 32000)
660 | # c2 = np.random.randn(4, 32000)
661 | # print(f'e1.shape {e1.shape} e2.shape {e2.shape}')
662 | # print(f'c1.shape {c1.shape} c2.shape {c2.shape}')
663 | # print(pit_sisdr_numpy(e1, e2, c1, c2))
664 | inputs = torch.randn((4, 6, 64000))
665 | model = Tree()
666 | model.to(torch.device('cpu'))
667 | sps_1, sps_2, es_sig_1, es_sig_2 = model(inputs)
668 | print(f'sps_1.shape {sps_1.shape} sps_2.shape {sps_2.shape}', flush=True)
669 | print(f'es_sig_1.shape {es_sig_1.shape} es_sig_2.shape {es_sig_2.shape}', flush=True)
670 |
671 |
672 |
673 |
--------------------------------------------------------------------------------
/preparation/1_sb_vad.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import sys
3 | sys.path.append("/Work21/2021/fuyanjie/libs/speechbrain/speechbrain/pretrained")
4 | from interfaces import VAD
5 | import os
6 | import glob
7 |
8 | """
9 | RUN this script for trainset, devset and testset, respectively
10 | """
11 |
12 | device = torch.device('cuda')
13 | # device = torch.device('cpu')
14 |
15 | VAD = VAD.from_hparams(source="speechbrain/vad-crdnn-libriparty", savedir="pretrained_models/vad-crdnn-libriparty")
16 | VAD.to(device)
17 |
18 | # wav_path_pattern = '/CDShare3/LibriSpeech/train-clean-100/*/*/*.flac'
19 | wav_path_pattern = '/CDShare3/LibriSpeech/dev-clean/*/*/*.flac'
20 | files = glob.glob(wav_path_pattern)
21 |
22 | # output_folder = '/Work21/2021/fuyanjie/pycode/LaBNet/data/metadata/vad/train-clean-100/'
23 | output_folder = '/Work21/2021/fuyanjie/pycode/LaBNet/data/metadata/vad/dev-clean/'
24 |
25 | for audio_file in files:
26 | try:
27 | print(f'{audio_file}', flush=True)
28 | boundaries = VAD.get_speech_segments(audio_file)
29 |
30 | file_name = str(audio_file).split('/')[-1]
31 | spkid = file_name.split('-')[0]
32 |
33 | save_folder = os.path.join(output_folder, spkid)
34 | if not os.path.exists(save_folder):
35 | os.makedirs(save_folder)
36 | save_path = os.path.join(save_folder, file_name.replace('.flac', '.txt'))
37 | print('save_path ', save_path, flush=True)
38 | # Print the output
39 | VAD.save_boundaries(boundaries, save_path=save_path)
40 | except Exception as e:
41 | print(f'Exception: {e}')
42 | continue
--------------------------------------------------------------------------------
/preparation/2_clip_segments.py:
--------------------------------------------------------------------------------
1 | import os
2 | import numpy as np
3 | import soundfile as sf
4 |
5 | """
6 | RUN this script for trainset, devset and testset, respectively
7 | """
8 |
9 | SR = 16000
10 | SEG_DURATION = 4 # duration in seconds
11 | SEG_LEN = SEG_DURATION * SR
12 |
13 | def clip_train_dev_test_segs(dataset_split):
14 | librispeech_path = f"/CDShare3/LibriSpeech/{dataset_split}"
15 | vad_results_path = f"/Work21/2021/fuyanjie/pycode/LaBNet/data/metadata/vad/{dataset_split}"
16 | output_libri_segments_dir = f"/CDShare3/LibriSpeechSegments/{dataset_split}"
17 |
18 | vad_res_spks = os.listdir(vad_results_path)
19 | for vad_res_spk in vad_res_spks:
20 | spkid = vad_res_spk
21 | vad_res_spk = os.path.join(vad_results_path, vad_res_spk)
22 | vad_res_utts = os.listdir(vad_res_spk)
23 | for vad_res_utt in vad_res_utts:
24 | uttid = vad_res_utt[:-4]
25 | chapid = uttid.split('-')[1]
26 | vad_res_utt = os.path.join(vad_res_spk, vad_res_utt)
27 | lines = open(vad_res_utt, 'r').readlines()
28 | start_idx = None
29 | end_idx = None
30 | if len(lines) == 1:
31 | seg_data = lines[0].split(' ')
32 | if float(seg_data[4]) < SEG_DURATION or seg_data[-1].startswith('N'):
33 | continue
34 | print(f'seg_data {seg_data}', flush=True)
35 | start_idx = 0
36 | end_idx = float(seg_data[4]) * SR
37 | else:
38 | for row_idx, line in enumerate(lines):
39 | seg_data = line.split(' ')
40 | print(f'seg_data {seg_data}', flush=True)
41 | if seg_data[-1].startswith('S'):
42 | duration = float(seg_data[4]) - float(seg_data[2])
43 | if duration > 4:
44 | start_idx = float(seg_data[2]) * SR
45 | end_idx = float(seg_data[4]) * SR
46 | if not start_idx or not end_idx:
47 | continue
48 |
49 | utt_path = os.path.join(librispeech_path, spkid, chapid, uttid + '.flac')
50 | wave_data, sr = sf.read(utt_path)
51 | end_idx = int(start_idx) + SEG_LEN
52 | seg_data = wave_data[int(start_idx):end_idx]
53 | dst_parent_dir = os.path.join(output_libri_segments_dir, spkid, chapid)
54 | os.makedirs(dst_parent_dir, exist_ok=True)
55 | dst_path = os.path.join(dst_parent_dir, uttid + '.flac')
56 | print(f'dst_path {dst_path}', flush=True)
57 | sf.write(dst_path, seg_data, 16000)
58 |
59 | def clip_wer_test_segs(dataset_split):
60 | librispeech_path = f"/CDShare3/LibriSpeech/{dataset_split}"
61 | vad_results_path = f"/Work21/2021/fuyanjie/pycode/LaBNet/data/metadata/vad/{dataset_split}"
62 | output_libri_segments_dir = f"/CDShare3/LibriSpeechSegments/{dataset_split}_wer"
63 |
64 | vad_res_spks = os.listdir(vad_results_path)
65 | vad_res_spks.sort()
66 | index = 0
67 | for vad_res_spk in vad_res_spks:
68 | spkid = vad_res_spk
69 | vad_res_spk = os.path.join(vad_results_path, vad_res_spk)
70 | vad_res_utts = os.listdir(vad_res_spk)
71 | vad_res_utts.sort()
72 | for vad_res_utt in vad_res_utts:
73 | try:
74 | index += 1
75 | random_state = np.random.RandomState(index)
76 |
77 | uttid = vad_res_utt[:-4]
78 | chapid = uttid.split('-')[1]
79 | vad_res_utt = os.path.join(vad_res_spk, vad_res_utt)
80 | lines = open(vad_res_utt, 'r').readlines()
81 | print(f'vad_res_utt {vad_res_utt}', flush=True)
82 |
83 | seg_data = lines[-1].split(' ')
84 | if float(seg_data[4]) > SEG_DURATION:
85 | continue
86 | print(f'seg_data {seg_data}', flush=True)
87 |
88 | utt_path = os.path.join(librispeech_path, spkid, chapid, uttid + '.flac')
89 | wave_data, sr = sf.read(utt_path)
90 | # print(f'{wave_data.shape} {wave_data.shape[-1]}')
91 | if wave_data.shape[-1] > SEG_LEN or wave_data.shape[-1] < SR * 3.5:
92 | continue
93 | # randomly padding to 4 secs long segment
94 | wave_padding_data = np.zeros(SEG_LEN)
95 |
96 | start_idx = random_state.randint(0, SEG_LEN - wave_data.shape[-1])
97 | end_idx = start_idx + len(wave_data)
98 | wave_padding_data[start_idx:end_idx] = wave_data
99 | seg_data = wave_padding_data
100 | print(f'start_idx {start_idx} end_idx {end_idx}', flush=True)
101 |
102 | dst_parent_dir = os.path.join(output_libri_segments_dir, spkid, chapid)
103 | os.makedirs(dst_parent_dir, exist_ok=True)
104 | dst_path = os.path.join(dst_parent_dir, uttid + '.flac')
105 | print(f'dst_path {dst_path}', flush=True)
106 | sf.write(dst_path, seg_data, 16000)
107 | except Exception as e:
108 | print(f'Exception: {e}')
109 | continue
110 |
111 | if __name__ == '__main__':
112 | # dataset_split = 'train-clean-100'
113 | # clip_train_dev_test_segs(dataset_split)
114 | # dataset_split = 'dev-clean'
115 | # clip_train_dev_test_segs(dataset_split)
116 | # dataset_split = 'test-clean'
117 | # clip_train_dev_test_segs(dataset_split)
118 | dataset_split = 'test-clean'
119 | clip_wer_test_segs(dataset_split)
120 |
--------------------------------------------------------------------------------
/preparation/3_generate_list.py:
--------------------------------------------------------------------------------
1 | import librosa
2 | import os
3 | import math
4 | import soundfile as sf
5 |
6 | def generate_flac_list(libri_segments_dir, output_list_path):
7 | spkids = os.listdir(libri_segments_dir)
8 | count = 0
9 | with open(output_list_path, 'w', encoding='utf-8') as output_f:
10 | for spkid in spkids:
11 | spk_dir = os.path.join(libri_segments_dir, spkid)
12 | chapids = os.listdir(spk_dir)
13 | for chapid in chapids:
14 | chap_dir = os.path.join(spk_dir, chapid)
15 | wav_names = os.listdir(chap_dir)
16 | for wav_name in wav_names:
17 | wav_path = os.path.join(chap_dir, wav_name)
18 | print(f'wav_path {wav_path}')
19 | output_f.write(wav_path + '\n')
20 | count += 1
21 | print('Finish count:{} | {}'.format(count, wav_path))
22 |
23 | def generate_spk_list(libri_segments_dir, output_list_path):
24 | spkids = os.listdir(libri_segments_dir)
25 | spkids.sort()
26 | with open(output_list_path, 'w', encoding='utf-8') as output_f:
27 | for spkid in spkids:
28 | output_f.write(spkid + '\n')
29 |
30 | def summarize_transcription(librispeech_path, libri_segments_dir, output_txt_path):
31 | spkids = os.listdir(libri_segments_dir)
32 | count = 0
33 | with open(output_txt_path, 'w', encoding='utf-8') as output_f:
34 | for spkid in spkids:
35 | spk_dir = os.path.join(libri_segments_dir, spkid)
36 | chapids = os.listdir(spk_dir)
37 | for chapid in chapids:
38 | chap_dir = os.path.join(spk_dir, chapid)
39 | wav_names = os.listdir(chap_dir)
40 | uttids = []
41 | for wav_name in wav_names:
42 | uttid = wav_name[:-5]
43 | print(f'uttid {uttid}')
44 | uttids.append(uttid)
45 | src_txt_path = os.path.join(librispeech_path, spkid, chapid, f'{spkid}-{chapid}.trans.txt')
46 | src_txt = open(src_txt_path, 'r').readlines()
47 | for line in src_txt:
48 | line = line.strip()
49 | uttid, txt = line.split(' ', 1)
50 | print(f'uttid: {uttid} txt: {txt}')
51 | if uttid in uttids:
52 | output_f.write(uttid + ' ' + txt + '\n')
53 | count += 1
54 | print('Finish count:{} | {}'.format(count, uttid))
55 |
56 |
57 | if __name__ == '__main__':
58 | # dataset_split = 'test-clean'
59 | dataset_split = 'test-clean_wer'
60 | libri_segments_dir = f'/CDShare3/LibriSpeechSegments/{dataset_split}'
61 | output_flac_list_path = f'/Work21/2021/fuyanjie/pycode/LaBNet/data/libri_segments_list/{dataset_split}.lst'
62 | output_spk_list_path = f'/Work21/2021/fuyanjie/pycode/LaBNet/data/libri_segments_list/{dataset_split}_spk.lst'
63 | generate_flac_list(libri_segments_dir, output_flac_list_path)
64 | generate_spk_list(libri_segments_dir, output_spk_list_path)
65 |
66 | librispeech_path = f"/CDShare3/LibriSpeech/test-clean"
67 | libri_segments_dir = '/CDShare3/LibriSpeechSegments/test-clean_wer'
68 | output_txt_path = '/CDShare3/LibriSpeechSegments/test-clean_wer/test-clean_wer.txt'
69 | summarize_transcription(librispeech_path, libri_segments_dir, output_txt_path)
--------------------------------------------------------------------------------
/preparation/libri_sim_code/4_gen_simulated_room.py:
--------------------------------------------------------------------------------
1 | import math
2 | import multiprocessing as mp
3 | import os
4 | import random
5 | import re
6 | import shutil
7 | import time
8 | from concurrent.futures.process import ProcessPoolExecutor
9 | from functools import partial
10 | from joblib import Parallel, delayed
11 | import sys
12 | import wave
13 | import glob
14 | sys.path.append("/Work21/2021/fuyanjie/pycode/LaBNet/libri_sim_code")
15 |
16 | import numpy as np
17 | import soundfile as sf
18 |
19 | from gen_room_para import gen_room_para, gen_mulchannel_data_random, gen_mulchannel_data_angle
20 |
21 | MAX_WORKERS = 10
22 | SEED = 42
23 |
24 | def get_first_level_folder(dir):
25 | folder_list = []
26 | for entry in os.scandir(dir):
27 | if entry.is_dir():
28 | folder_list.append(entry.path)
29 |
30 | return folder_list
31 |
32 |
33 | def generate_wav_list_from_lst_file(lst_file):
34 | with open(lst_file, "r", encoding='utf-8-sig') as f:
35 | file_list = f.readlines()
36 |
37 | return file_list
38 |
39 |
40 | def save_audio_separately(audio, path, fs):
41 | channels, length = audio.shape
42 |
43 | for channel in range(channels):
44 | suffix = '_multichannel_' + str(channel) + '.wav'
45 |
46 | out_data = np.reshape(audio[channel, :], [length, 1])
47 |
48 | out_data = out_data.astype(np.int16)
49 |
50 | with wave.open(path+suffix, 'wb') as f:
51 | f.setframerate(fs)
52 | f.setsampwidth(2)
53 | f.setnchannels(1)
54 | f.writeframes(out_data.tostring())
55 |
56 |
57 | def __sim_stationary_noise(folder, max_gen_num, wav_list, fs, speech_length):
58 | room_para_path = os.path.join(folder, 'room_para.npy')
59 | room_para = np.load(room_para_path, allow_pickle=True).item()
60 |
61 | # used_wav_num = random.randint(5, 10)
62 | used_wav_num = 1
63 | wav_samples = random.sample(wav_list, used_wav_num)
64 | print(f'num of wav_samples {len(wav_samples)}')
65 |
66 | # judge if multi-channel stationary noise wav files exist
67 | # file_list = os.listdir(os.path.join(folder, 'stationary_noise'))
68 | st_noise_dir = os.path.join(folder, 'stationary_noise')
69 | if os.path.exists(st_noise_dir) and len(os.listdir(st_noise_dir)) > 0:
70 | print(f'{st_noise_dir} is not empty!')
71 | return
72 |
73 | for num in range(max_gen_num):
74 | # 生成多个多通道平稳噪声数据
75 | mulchannel_audio_data_list = []
76 | for wav_file in wav_samples:
77 | multichannel_audio_data, _, angle_degree = gen_mulchannel_data_random(wav_file.strip(),
78 | room_para,
79 | folder,
80 | audio_type=0,
81 | fs=fs,
82 | segment_length=speech_length)
83 | multichannel_audio_data = multichannel_audio_data[:, 0:fs * speech_length]
84 | mulchannel_audio_data_list.append(multichannel_audio_data)
85 |
86 | # 对生成数据幅度进行一下调整，再进行叠加
87 | out_mulchannel_audio_data = np.zeros_like(mulchannel_audio_data_list[0], dtype=np.int16)
88 | for data in mulchannel_audio_data_list:
89 | out_mulchannel_audio_data += data
90 |
91 | # 将生成好的数据进行保存
92 | mic_num, _ = out_mulchannel_audio_data.shape
93 | save_folder = os.path.join(folder, 'stationary_noise')
94 | output_file_path = os.path.join(save_folder, 'stationary_noise_{}'.format(num))
95 | save_audio_separately(out_mulchannel_audio_data, output_file_path, fs=fs)
96 |
97 |
98 | def sim_stationary_noise(base_dir, stationary_noise_lst_file, sim_room_num, max_stationary_noise_num, fs, room_transpose_prob, speech_length):
99 | '''
100 | 利用多个平稳噪声模拟真实环境下的多通道平稳噪声，每个房间就仿真一条数据，主要目的是为了更好模拟不同环境下的平稳噪声
101 | :param base_dir: 根目录，用于保存不同房间的数据
102 | :param stationary_noise_lst_file: 平稳噪声列表文件
103 | :param sim_room_num: 模拟不同房间的数目
104 | :param max_stationary_noise_num: 最多使用平稳噪声的数目
105 | :param fs:
106 | :return:
107 | '''
108 | for i in range(sim_room_num):
109 | _ = gen_room_para(base_dir, room_transpose_prob=room_transpose_prob)
110 |
111 | room_folders = get_first_level_folder(base_dir)
112 | wav_list = list(generate_wav_list_from_lst_file(stationary_noise_lst_file))
113 |
114 | Parallel(n_jobs=MAX_WORKERS)(delayed(__sim_stationary_noise)(room_folder, max_stationary_noise_num, wav_list, fs, speech_length) for room_folder in room_folders)
115 | # with ProcessPoolExecutor(MAX_WORKERS) as ex:
116 | # func = partial(__sim_stationary_noise, max_gen_num=max_stationary_noise_num, wav_list=wav_list,
117 | # fs=fs, speech_length=speech_length)
118 | # ex.map(func, room_folders)
119 |
120 |
121 | def __sim_speech(seed_idx, folder, libri_segments_dir, fs, cover_angle_range, multi_list, segment_length, speaker_list, angular_spacing):
122 | random_state = np.random.RandomState(seed_idx * SEED)
123 |
124 | room_para_path = os.path.join(folder, 'room_para.npy')
125 | room_para = np.load(room_para_path, allow_pickle=True).item()
126 |
127 | # if room_para['room_dim'][1] < 5:
128 | # cover_angle_range = [20, 160]
129 | # elif room_para['room_dim'][1] >= 5 and room_para['room_dim'][1] < 7:
130 | # cover_angle_range = [15, 165]
131 | # elif room_para['room_dim'][1] >= 7:
132 | # cover_angle_range = [10, 170]
133 | # 声明多通道人声数据保存路径
134 | mulchannel_speech_folder = os.path.join(folder, 'speech')
135 |
136 | cover_angle_num = (cover_angle_range[1] - cover_angle_range[0]) // angular_spacing
137 |
138 | speaker_folder = ""
139 | speaker_folders = []
140 | speaker_lsts = []
141 | seen_speaker_lsts = []
142 |
143 | with open(speaker_list) as fid:
144 | for line in fid:
145 | seen_speaker_lsts.append(line.strip())
146 |
147 | if not multi_list:
148 | speaker_folder = libri_segments_dir
149 | print(f'speaker_folder {speaker_folder}', flush=True)
150 | speaker_lst = seen_speaker_lsts
151 | print(f'speaker_lst {speaker_lst}', flush=True)
152 | speaker_idx_lst = np.arange(len(speaker_lst))
153 | random_state.shuffle(speaker_idx_lst)
154 | speaker_idx_lst = speaker_idx_lst[:cover_angle_num * 2]
155 |
156 | # get the wav_id list
157 | wav_id_file_path = os.path.join(folder, 'speech_wav_id_lst.lst')
158 | if os.path.exists(wav_id_file_path):
159 | wav_id_lst = []
160 | with open(wav_id_file_path, 'r', encoding='utf-8') as f:
161 | wav_id_set = f.readlines()
162 |
163 | for wav_id in wav_id_set:
164 | wav_id = wav_id.strip()
165 | wav_id_lst.append(wav_id)
166 | else:
167 | wav_id_lst = []
168 | angle = cover_angle_range[0]
169 | index = 0
170 | while True:
171 | index %= len(speaker_idx_lst)
172 | k = speaker_idx_lst[index]
173 | index += 1
174 | wav_lst = glob.glob(os.path.join(speaker_folder, speaker_lst[k]) + '/*/*.flac')
175 | wav_name = random.sample(wav_lst, 1)[0]
176 |
177 | if speaker_lst[k] not in wav_id_lst:
178 | wav_path = os.path.join(speaker_folder, speaker_lst[k])
179 | wav_path = os.path.join(wav_name.split('-')[1], wav_name)
180 |
181 | # if angle >= 55:
182 | # angular_spacing = 5
183 | # if angle >= 120:
184 | # angular_spacing = 10
185 | # random_range = random.randint(0, angular_spacing - 1)
186 | # random_angle = angle + random_range
187 | random_angle = angle
188 |
189 |
190 | for dis_idx in range(3):
191 | multichannel_audio_data, wav_id, distance = gen_mulchannel_data_angle(wav_path,
192 | room_para,
193 | folder,
194 | angle=random_angle,
195 | distance_flag=dis_idx,
196 | fs=fs,
197 | segment_length=segment_length,
198 | audio_type=1)
199 |
200 | max_speech_value = np.max(np.abs(multichannel_audio_data))
201 |
202 |
203 | if max_speech_value > 32767:
204 | multichannel_audio_data = multichannel_audio_data / max_speech_value * 30000
205 | out_mulchannel_audio_data = multichannel_audio_data
206 |
207 |
208 | # 生成文件保存路径
209 | if dis_idx == 0:
210 | distance_type = "close"
211 | elif dis_idx == 1:
212 | distance_type = "middle"
213 | else:
214 | distance_type = "far"
215 | file_name = "{}-{}m-{}".format(wav_id, distance, distance_type)
216 | angle_speech_folder = os.path.join(mulchannel_speech_folder, f"{random_angle}-{wav_id.split('-')[0]}")
217 | os.makedirs(angle_speech_folder, exist_ok=True)
218 | save_file_path = os.path.join(angle_speech_folder, file_name)
219 |
220 | save_audio_separately(out_mulchannel_audio_data, save_file_path, fs=fs)
221 | print(save_file_path + " finish!")
222 |
223 | wav_id_lst.append(wav_id)
224 | angle = angle + angular_spacing
225 | if angle >= cover_angle_range[1] + 1 - angular_spacing:
226 | break
227 | else:
228 | continue
229 | # wait for all the wav simulation finish, write wav_id_lst in the wav_id_file_path
230 | with open(wav_id_file_path, 'w', encoding='utf-8') as f:
231 | for wav_id in wav_id_lst:
232 | f.write(wav_id + '\n')
233 |
234 | def sim_speech(base_dir, libri_segments_dir, fs, cover_angle_range, segment_length, speaker_list,
235 | angular_spacing):
236 | '''
237 | 模拟多通道人声数据
238 | :param base_dir: 根目录
239 | :param libri_segments_dir: 单通道人声文件夹
240 | :param max_wav_num:
241 | :param fs:
242 | :return:
243 | '''
244 | room_folders = get_first_level_folder(base_dir)
245 |
246 | Parallel(n_jobs=MAX_WORKERS)(delayed(__sim_speech)(seed_idx, folder, libri_segments_dir, fs, cover_angle_range, False, segment_length, speaker_list, angular_spacing) for seed_idx, folder in enumerate(room_folders))
247 |
248 | # with ProcessPoolExecutor(MAX_WORKERS) as ex:
249 | # func = partial(__sim_speech, libri_segments_dir=libri_segments_dir, fs=fs, cover_angle_range=cover_angle_range, multi_list=False,
250 | # segment_length=segment_length, speaker_list=speaker_list, angular_spacing=angular_spacing)
251 | # ex.map(func, room_folders)
252 | # debug
253 | # for folder in room_folders:
254 | # __sim_speech(folder, libri_segments_dir, fs, cover_angle_range, False, segment_length, speaker_list, angular_spacing)
255 | # debug
256 |
257 |
258 | def __sim_non_stationary_noise(folder, sim_wav_num, wav_list, fs, cover_angle_range, speech_length, angular_spacing):
259 | room_para_path = os.path.join(folder, 'room_para.npy')
260 | room_para = np.load(room_para_path, allow_pickle=True).item()
261 |
262 | speech_path = os.path.join(folder, 'speech')
263 | angle_of_speech = os.listdir(speech_path)
264 | angle_of_speech = list(map(int, angle_of_speech))
265 | angle_of_speech.sort()
266 | print(f'angle_list_of_speech: {angle_of_speech}')
267 |
268 | # 声明多通道非平稳噪声数据保存路径
269 | mulchannel_nonstationary_noise_folder = os.path.join(folder, 'nonstationary_noise')
270 |
271 | cover_angle_num = (cover_angle_range[1] - cover_angle_range[0]) / angular_spacing
272 | cover_angle_num = math.floor(cover_angle_num)
273 |
274 | non_stationary_list = wav_list
275 | np.random.shuffle(non_stationary_list)
276 | non_stationary_list = non_stationary_list[:3 * cover_angle_num]
277 |
278 | angle = angle_of_speech[0]
279 | index = 0
280 | for wav in non_stationary_list:
281 | multichannel_audio_data, wav_id, distance = gen_mulchannel_data_angle(wav.strip(),
282 | room_para,
283 | folder,
284 | angle=angle,
285 | distance_flag=index % 3,
286 | fs=fs,
287 | segment_length=speech_length,
288 | audio_type=2)
289 |
290 | # out_mulchannel_audio_data = multichannel_audio_data - np.mean(multichannel_audio_data, axis=0)
291 | out_mulchannel_audio_data = multichannel_audio_data
292 |
293 | # 生成文件保存路径
294 | if index % 3 == 0:
295 | distance_type = "close"
296 | elif index % 3 == 1:
297 | distance_type = "middle"
298 | else:
299 | distance_type = "far"
300 | file_name = "{}-{}m-{}".format(wav_id, distance, distance_type)
301 | angle_speech_folder = os.path.join(mulchannel_nonstationary_noise_folder, str(angle))
302 |
303 | save_file_path = os.path.join(angle_speech_folder, file_name)
304 |
305 | save_audio_separately(out_mulchannel_audio_data, save_file_path, fs=fs)
306 | print(save_file_path + " finish!")
307 |
308 | if index % 3 == 2:
309 | angle = angle + angular_spacing
310 | index = index + 1
311 |
312 |
313 |
314 | def sim_non_stationary_noise(base_dir, mono_nonstationary_noise_lst_file, max_wav_num, fs, cover_angle_range,
315 | speech_length, angular_spacing):
316 | '''
317 | 模拟多通道非平稳噪声
318 | :param base_dir: 根目录
319 | :param mono_nonstationary_noise_lst_file: 单通道非平稳噪声文件列表文件
320 | :param max_wav_num:
321 | :param fs: 采样率
322 | :return:
323 | '''
324 | room_folders = get_first_level_folder(base_dir)
325 |
326 | wav_list = list(generate_wav_list_from_lst_file(mono_nonstationary_noise_lst_file))
327 | if len(wav_list) > max_wav_num:
328 | sim_wav_num = max_wav_num
329 | else:
330 | sim_wav_num = len(wav_list)
331 |
332 | with ProcessPoolExecutor(MAX_WORKERS) as ex:
333 | func = partial(__sim_non_stationary_noise, sim_wav_num=sim_wav_num, wav_list=wav_list, fs=fs,
334 | cover_angle_range=cover_angle_range, speech_length=speech_length,
335 | angular_spacing=angular_spacing)
336 | ex.map(func, room_folders)
337 |
338 |
339 | def single_proc(stationary_noise_lst, libri_segments_dir, non_stationary_noise_lst, save_path,
340 | sim_room_num, stationary_noise_num, fs, speech_length, cover_angle_range, room_transpose_prob, speaker_list, augular_spacing):
341 | print("MAX_WORKERS:", MAX_WORKERS)
342 |
343 | # Start timing¬
344 | start_time = time.time()
345 |
346 | # make directory of output data
347 | if not os.path.exists(save_path):
348 | os.makedirs(save_path)
349 | else:
350 | print(f'{save_path} already exists!')
351 | shutil.rmtree(save_path)
352 | print(f'remove {save_path}')
353 |
354 | # generate room parameters and stationary noise
355 | sim_stationary_noise(save_path, stationary_noise_lst, sim_room_num, stationary_noise_num, fs, room_transpose_prob, speech_length)
356 | print('Finish sim_stationary_noise')
357 |
358 | # generate multi-channel speech
359 | sim_speech(save_path, libri_segments_dir, fs, cover_angle_range, speech_length, speaker_list, augular_spacing)
360 | print('Finish sim_speech')
361 |
362 | # DO NOT INVOKE FOR NOW!
363 | # generate multi-channel non-stationary noise
364 | # sim_non_stationary_noise(save_path, non_stationary_noise_lst, non_stationary_num, fs, cover_angle_range, speech_length, augular_spacing)
365 | # print('Finish sim_non_stationary_noise')
366 |
367 | # End timing
368 | print('time spent: %s mins' % str((time.time() - start_time) / 60.0))
369 |
370 |
371 | if __name__ == '__main__':
372 | '''
373 | convert single channel format to multi-channel(6-channel) format of single voice.
374 |
375 | :param stationary_noise_lst: contains the absolute path of stationary noise .wav
376 | :param libri_segments_dir: contains clean audio files
377 | :param non_stationary_noise_lst: contains the absolute path of non-stationary noise .wav
378 | :param save_path: the output path of multi-channel data
379 | :param sim_room_num: the number of simulated rooms, each room contains specified number of stationary noise, speech, and non-stationary noise
380 | :param stationary_noise_num: the upper bound of the kind of stationary noise we choose
381 | :param speech_num: the number of simulated speech in each room
382 | :param non_stationary_num: the number of simulated non-stationary noise in each room
383 | :param fs: sample rate
384 | :return:
385 | '''
386 |
387 | # 设置路径参数
388 | # dataset_split = "train-clean-100"
389 | # dataset_split = "dev-clean"
390 | # dataset_split = "test-clean"
391 | dataset_split = "test-clean_wer"
392 | stationary_noise_lst = '/Work21/2020/yinhaoran/VCTK_simulated_data/list/clean_speech_3.lst'
393 | libri_segments_dir = f'/CDShare3/LibriSpeechSegments/{dataset_split}'
394 | non_stationary_noise_lst = '/non_stationary_noise.lst'
395 |
396 | speaker_lst = f'/Work21/2021/fuyanjie/pycode/LaBNet/data/libri_segments_list/{dataset_split}_spk.lst'
397 |
398 | # 设定合成参数
399 | sim_room_num = 10
400 | stationary_noise_num = 1
401 | fs = 16000
402 | speech_length = 4
403 | cover_angle_range = [0, 180]
404 | angular_spacing = 1 # 角度最小间隔
405 | room_transpose_prob = 0.5
406 |
407 | save_path = f'/CDShare3/Libri-SIM/rooms_0129/{dataset_split}_{sim_room_num}rooms'
408 |
409 | single_proc(stationary_noise_lst, libri_segments_dir, non_stationary_noise_lst, save_path, sim_room_num,
410 | stationary_noise_num, fs, speech_length, cover_angle_range, room_transpose_prob,
411 | speaker_lst, angular_spacing)
412 |
413 |
--------------------------------------------------------------------------------
/preparation/libri_sim_code/5_generate_json.py:
--------------------------------------------------------------------------------
1 | # -*- coding:utf-8 _*-
2 | from itertools import groupby
3 | import json
4 | import os
5 | import pickle
6 | import random
7 | from acoustics.signal import highpass
8 |
9 | import librosa
10 | import numpy as np
11 | import soundfile as sf
12 |
13 | SEED = 7
14 | random.seed(SEED)
15 | sr = 16000
16 |
17 | def audioread(path, fs=16000):
18 | wave_data, sr = sf.read(path)
19 | if sr != fs:
20 | wave_data = librosa.resample(wave_data, sr, fs)
21 | return wave_data
22 |
23 |
24 | def generate_vad_label(orig_signal, sr = 16000, win_len = 0.032, hop_len = 0.016):
25 | original_signal = orig_signal / np.max(orig_signal) * 25000
26 | original_signal = highpass(original_signal, 100, sr, order =8)
27 |
28 | vad = np.ones(int(len(original_signal)/(hop_len * sr)))
29 |
30 | frame_shift = int(hop_len * sr) # 256
31 | frame_len = int(win_len * sr) # 512
32 |
33 | value_threshold = 110
34 |
35 | sig_100ms_len = int(0.01 * sr)
36 | frame_max_value = np.max(original_signal[:sig_100ms_len])
37 | if frame_max_value > value_threshold:
38 | value_threshold = frame_max_value
39 |
40 | for i in range(len(vad)):
41 | frame_data = original_signal[i * frame_shift : (i + 1) * frame_shift] # old: i -> i+1, changed： i -> i + length
42 | frame_data = np.abs(frame_data)
43 | frame_data = np.where(frame_data > value_threshold,0,1)
44 | if np.sum(frame_data) > frame_shift * 0.9 and vad[i] == 1:
45 | vad[i] = 0
46 |
47 | mid_len_speech = round(0.1 / hop_len + 0.1)
48 | mid_len_non_speech = round(0.05 / hop_len + 0.1)
49 | for i in range(len(vad)):
50 | if i < len(vad) - mid_len_speech - 2:
51 | if vad[i] == 0 and vad[i+mid_len_speech+1] == 0 and sum(vad[i+1:i+mid_len_speech+1]
48 |