Spaces:
Running
Running
File size: 2,796 Bytes
406f22d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import csv
import torch
import numpy as np
import logging
from torch_mir_eval.separation import bss_eval_sources
import fast_bss_eval
from ..losses import (
PITLossWrapper,
pairwise_neg_sisdr,
pairwise_neg_snr,
singlesrc_neg_sisdr,
PairwiseNegSDR,
)
logger = logging.getLogger(__name__)
class MetricsTracker:
def __init__(self, save_file: str = ""):
self.all_sdrs = []
self.all_sdrs_i = []
self.all_sisnrs = []
self.all_sisnrs_i = []
csv_columns = ["snt_id", "sdr", "sdr_i", "si-snr", "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(
PairwiseNegSDR("sisdr", zero_mean=False), pit_from="pw_mtx"
)
self.pit_snr = PITLossWrapper(
PairwiseNegSDR("snr", zero_mean=False), pit_from="pw_mtx"
)
def __call__(self, mix, clean, estimate, key):
# sisnr
sisnr = self.pit_sisnr(estimate.unsqueeze(0), clean.unsqueeze(0))
mix = torch.stack([mix] * clean.shape[0], dim=0)
sisnr_baseline = self.pit_sisnr(mix.unsqueeze(0), clean.unsqueeze(0))
sisnr_i = sisnr - sisnr_baseline
# sdr
sdr = -fast_bss_eval.sdr_pit_loss(estimate, clean).mean()
sdr_baseline = -fast_bss_eval.sdr_pit_loss(mix, clean).mean()
sdr_i = sdr - sdr_baseline
# import pdb; pdb.set_trace()
row = {
"snt_id": key,
"sdr": sdr.item(),
"sdr_i": sdr_i.item(),
"si-snr": -sisnr.item(),
"si-snr_i": -sisnr_i.item(),
}
self.writer.writerow(row)
# Metric Accumulation
self.all_sdrs.append(sdr.item())
self.all_sdrs_i.append(sdr_i.item())
self.all_sisnrs.append(-sisnr.item())
self.all_sisnrs_i.append(-sisnr_i.item())
def update(self, ):
return {"sdr_i": np.array(self.all_sdrs_i).mean(),
"si-snr_i": np.array(self.all_sisnrs_i).mean()
}
def final(self,):
row = {
"snt_id": "avg",
"sdr": np.array(self.all_sdrs).mean(),
"sdr_i": np.array(self.all_sdrs_i).mean(),
"si-snr": np.array(self.all_sisnrs).mean(),
"si-snr_i": np.array(self.all_sisnrs_i).mean(),
}
self.writer.writerow(row)
row = {
"snt_id": "std",
"sdr": np.array(self.all_sdrs).std(),
"sdr_i": np.array(self.all_sdrs_i).std(),
"si-snr": np.array(self.all_sisnrs).std(),
"si-snr_i": np.array(self.all_sisnrs_i).std(),
}
self.writer.writerow(row)
self.results_csv.close()
|