|
import torch |
|
from torchaudio import load |
|
|
|
from pesq import pesq |
|
from pystoi import stoi |
|
|
|
from .other import si_sdr, pad_spec |
|
|
|
|
|
sr = 16000 |
|
snr = 0.5 |
|
N = 30 |
|
corrector_steps = 1 |
|
|
|
|
|
def evaluate_model(model, num_eval_files): |
|
|
|
clean_files = model.data_module.valid_set.clean_files |
|
noisy_files = model.data_module.valid_set.noisy_files |
|
|
|
|
|
total_num_files = len(clean_files) |
|
indices = torch.linspace(0, total_num_files-1, num_eval_files, dtype=torch.int) |
|
clean_files = list(clean_files[i] for i in indices) |
|
noisy_files = list(noisy_files[i] for i in indices) |
|
|
|
_pesq = 0 |
|
_si_sdr = 0 |
|
_estoi = 0 |
|
|
|
for (clean_file, noisy_file) in zip(clean_files, noisy_files): |
|
|
|
x, _ = load(clean_file) |
|
y, _ = load(noisy_file) |
|
T_orig = x.size(1) |
|
|
|
|
|
norm_factor = y.abs().max() |
|
y = y / norm_factor |
|
|
|
|
|
Y = torch.unsqueeze(model._forward_transform(model._stft(y.cuda())), 0) |
|
Y = pad_spec(Y) |
|
y = y * norm_factor |
|
|
|
|
|
sampler = model.get_pc_sampler( |
|
'reverse_diffusion', 'ald', Y.cuda(), N=N, |
|
corrector_steps=corrector_steps, snr=snr) |
|
sample, _ = sampler() |
|
|
|
x_hat = model.to_audio(sample.squeeze(), T_orig) |
|
x_hat = x_hat * norm_factor |
|
|
|
x_hat = x_hat.squeeze().cpu().numpy() |
|
x = x.squeeze().cpu().numpy() |
|
y = y.squeeze().cpu().numpy() |
|
|
|
_si_sdr += si_sdr(x, x_hat) |
|
_pesq += pesq(sr, x, x_hat, 'wb') |
|
_estoi += stoi(x, x_hat, sr, extended=True) |
|
|
|
return _pesq/num_eval_files, _si_sdr/num_eval_files, _estoi/num_eval_files |
|
|
|
|
