Spaces:
Running
Running
| import csv | |
| import torch | |
| import numpy as np | |
| import logging | |
| # from torch_mir_eval.separation import bss_eval_sources | |
| from ..losses import ( | |
| PITLossWrapper, | |
| pairwise_neg_sisdr, | |
| pairwise_neg_snr, | |
| singlesrc_neg_sisdr, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class SPlitMetricsTracker: | |
| def __init__(self, save_file: str = ""): | |
| self.one_all_snrs = [] | |
| self.one_all_snrs_i = [] | |
| self.one_all_sisnrs = [] | |
| self.one_all_sisnrs_i = [] | |
| self.two_all_snrs = [] | |
| self.two_all_snrs_i = [] | |
| self.two_all_sisnrs = [] | |
| self.two_all_sisnrs_i = [] | |
| csv_columns = [ | |
| "snt_id", | |
| "one_snr", | |
| "one_snr_i", | |
| "one_si-snr", | |
| "one_si-snr_i", | |
| "two_snr", | |
| "two_snr_i", | |
| "two_si-snr", | |
| "two_si-snr_i", | |
| ] | |
| self.results_csv = open(save_file, "w") | |
| self.writer = csv.DictWriter(self.results_csv, fieldnames=csv_columns) | |
| self.writer.writeheader() | |
| self.pit_sisnr = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx") | |
| self.pit_snr = PITLossWrapper(pairwise_neg_snr, pit_from="pw_mtx") | |
| def __call__(self, mix, clean, estimate, key): | |
| _, ests_np = self.pit_snr( | |
| estimate.unsqueeze(0), clean.unsqueeze(0), return_ests=True | |
| ) | |
| # sisnr | |
| two_sisnr = self.pit_sisnr(ests_np[:, 0:2], clean.unsqueeze(0)[:, 0:2]) | |
| one_sisnr = self.pit_sisnr( | |
| ests_np[:, 2].unsqueeze(1), clean.unsqueeze(0)[:, 2].unsqueeze(1) | |
| ) | |
| mix = torch.stack([mix] * clean.shape[0], dim=0) | |
| two_sisnr_baseline = self.pit_sisnr( | |
| mix.unsqueeze(0)[:, 0:2], clean.unsqueeze(0)[:, 0:2] | |
| ) | |
| one_sisnr_baseline = self.pit_sisnr( | |
| mix.unsqueeze(0)[:, 2].unsqueeze(1), clean.unsqueeze(0)[:, 2].unsqueeze(1) | |
| ) | |
| two_sisnr_i = two_sisnr - two_sisnr_baseline | |
| one_sisnr_i = one_sisnr - one_sisnr_baseline | |
| # sdr | |
| two_snr = self.pit_snr(ests_np[:, 0:2], clean.unsqueeze(0)[:, 0:2]) | |
| one_snr = self.pit_snr( | |
| ests_np[:, 2].unsqueeze(1), clean.unsqueeze(0)[:, 2].unsqueeze(1) | |
| ) | |
| two_snr_baseline = self.pit_snr( | |
| mix.unsqueeze(0)[:, 0:2], clean.unsqueeze(0)[:, 0:2] | |
| ) | |
| one_snr_baseline = self.pit_snr( | |
| mix.unsqueeze(0)[:, 2].unsqueeze(1), clean.unsqueeze(0)[:, 2].unsqueeze(1) | |
| ) | |
| two_snr_i = two_snr - two_snr_baseline | |
| one_snr_i = one_snr - one_snr_baseline | |
| row = { | |
| "snt_id": key, | |
| "one_snr": -one_snr.item(), | |
| "one_snr_i": -one_snr_i.item(), | |
| "one_si-snr": -one_sisnr.item(), | |
| "one_si-snr_i": -one_sisnr_i.item(), | |
| "two_snr": -two_snr.item(), | |
| "two_snr_i": -two_snr_i.item(), | |
| "two_si-snr": -two_sisnr.item(), | |
| "two_si-snr_i": -two_sisnr_i.item(), | |
| } | |
| self.writer.writerow(row) | |
| # Metric Accumulation | |
| self.one_all_snrs.append(-one_snr.item()) | |
| self.one_all_snrs_i.append(-one_snr_i.item()) | |
| self.one_all_sisnrs.append(-one_sisnr.item()) | |
| self.one_all_sisnrs_i.append(-one_sisnr_i.item()) | |
| self.two_all_snrs.append(-two_snr.item()) | |
| self.two_all_snrs_i.append(-two_snr_i.item()) | |
| self.two_all_sisnrs.append(-two_sisnr.item()) | |
| self.two_all_sisnrs_i.append(-two_sisnr_i.item()) | |
| def final(self,): | |
| row = { | |
| "snt_id": "avg", | |
| "one_snr": np.array(self.one_all_snrs).mean(), | |
| "one_snr_i": np.array(self.one_all_snrs_i).mean(), | |
| "one_si-snr": np.array(self.one_all_sisnrs).mean(), | |
| "one_si-snr_i": np.array(self.one_all_sisnrs_i).mean(), | |
| "two_snr": np.array(self.two_all_snrs).mean(), | |
| "two_snr_i": np.array(self.two_all_snrs_i).mean(), | |
| "two_si-snr": np.array(self.two_all_sisnrs).mean(), | |
| "two_si-snr_i": np.array(self.two_all_sisnrs_i).mean(), | |
| } | |
| self.writer.writerow(row) | |
| # logger.info("Mean SISNR is {}".format(row["si-snr"])) | |
| # logger.info("Mean SISNRi is {}".format(row["si-snr_i"])) | |
| # logger.info("Mean SDR is {}".format(row["sdr"])) | |
| # logger.info("Mean SDRi is {}".format(row["sdr_i"])) | |
| self.results_csv.close() | |