| |
| |
| |
| |
| |
| |
| import csv |
| from sympy import im |
| import torch |
| import numpy as np |
| import logging |
| import os |
| import librosa |
| from torch_mir_eval.separation import bss_eval_sources |
| import fast_bss_eval |
| from visqol import visqol_lib_py |
| from visqol.pb2 import visqol_config_pb2 |
| from visqol.pb2 import similarity_result_pb2 |
|
|
| logger = logging.getLogger(__name__) |
|
|
| def is_silent(wav, threshold=1e-4): |
| return torch.sum(wav ** 2) / wav.numel() < threshold |
|
|
| class MetricsTracker: |
| def __init__(self, save_file: str = ""): |
| self.all_sdrs = [] |
| self.all_sisnrs = [] |
| self.all_visqols = [] |
| |
| csv_columns = ["snt_id", "sdr", "si-snr", "visqol"] |
| self.visqol_config = visqol_config_pb2.VisqolConfig() |
| self.visqol_config.audio.sample_rate = 48000 |
| self.visqol_config.options.use_speech_scoring = False |
| svr_model_path = "libsvm_nu_svr_model.txt" |
| self.visqol_config.options.svr_model_path = os.path.join(os.path.dirname(visqol_lib_py.__file__), "model", svr_model_path) |
| self.visqol_api = visqol_lib_py.VisqolApi() |
| self.visqol_api.Create(self.visqol_config) |
| |
| self.results_csv = open(save_file, "w") |
| self.writer = csv.DictWriter(self.results_csv, fieldnames=csv_columns) |
| self.writer.writeheader() |
| |
| def __call__(self, clean, estimate, key): |
| sisnr = fast_bss_eval.si_sdr(clean.unsqueeze(0), estimate.unsqueeze(0), zero_mean=True).mean() |
| sdr = fast_bss_eval.sdr(clean.unsqueeze(0), estimate.unsqueeze(0), zero_mean=True).mean() |
| |
| clean = librosa.resample(clean.squeeze(0).mean(0).cpu().numpy(), orig_sr=44100, target_sr=48000).astype(np.float64) |
| estimate = librosa.resample(estimate.squeeze(0).mean(0).cpu().numpy(), orig_sr=44100, target_sr=48000).astype(np.float64) |
| |
| visqol = self.visqol_api.Measure(clean, estimate).moslqo |
| |
| row = { |
| "snt_id": key, |
| "sdr": sdr.item(), |
| "si-snr": sisnr.item(), |
| "visqol": visqol |
| } |
| |
| self.writer.writerow(row) |
| |
| self.all_sdrs.append(sdr.item()) |
| self.all_sisnrs.append(sisnr.item()) |
| self.all_visqols.append(visqol) |
| |
| def update(self, ): |
| return {"sdr": np.array(self.all_sdrs).mean(), |
| "si-snr": np.array(self.all_sisnrs).mean(), |
| "visqol": np.array(self.all_visqols).mean()} |
| |
| def final(self,): |
| row = { |
| "snt_id": "avg", |
| "sdr": np.array(self.all_sdrs).mean(), |
| "si-snr": np.array(self.all_sisnrs).mean(), |
| "visqol": np.array(self.all_visqols).mean() |
| } |
| self.writer.writerow(row) |
| row = { |
| "snt_id": "std", |
| "sdr": np.array(self.all_sdrs).std(), |
| "si-snr": np.array(self.all_sisnrs).std(), |
| "visqol": np.array(self.all_visqols).std() |
| } |
| self.writer.writerow(row) |
| self.results_csv.close() |
|
|
