calc metrics also for 48 kHz

calc_metrics.py

@@ -5,6 +5,7 @@ from soundfile import read
 from tqdm import tqdm
 from pesq import pesq
 import pandas as pd
+import librosa
 
 from pystoi import stoi
 
@@ -19,21 +20,27 @@ if __name__ == '__main__':
     args = parser.parse_args()
 
     data = {"filename": [], "pesq": [], "estoi": [], "si_sdr": [], "si_sir": [],  "si_sar": []}
-    sr = 16000
 
     # Evaluate standard metrics
     noisy_files = []
     noisy_files += sorted(glob(join(args.noisy_dir, '*.wav')))
-    noisy_files += sorted(glob(join(args.test_dir, '**', '*.wav')))
+    noisy_files += sorted(glob(join(args.noisy_dir, '**', '*.wav')))
     for noisy_file in tqdm(noisy_files):
-        filename = noisy_file.split('/')[-1]
-        x, _ = read(join(args.clean_dir, filename))
-        y, _ = read(noisy_file)
+        filename = noisy_file.replace(args.noisy_dir, "")[1:]
+        if 'dB' in filename:
+            clean_filename = filename.split("_")[0] + ".wav"
+        else:
+            clean_filename = filename
+        x, sr_x = read(join(args.clean_dir, clean_filename))
+        y, sr_y = read(join(args.noisy_dir, filename))
+        x_hat, sr_x_hat = read(join(args.enhanced_dir, filename))
+        assert sr_x == sr_y == sr_x_hat
         n = y - x 
-        x_hat, _ = read(join(args.enhanced_dir, filename))
+        x_hat_16k = librosa.resample(x_hat, orig_sr=sr_x_hat, target_sr=16000) if sr_x_hat != 16000 else x_hat
+        x_16k = librosa.resample(x, orig_sr=sr_x, target_sr=16000) if sr_x != 16000 else x
         data["filename"].append(filename)
-        data["pesq"].append(pesq(sr, x, x_hat, 'wb'))
-        data["estoi"].append(stoi(x, x_hat, sr, extended=True))
+        data["pesq"].append(pesq(16000, x_16k, x_hat_16k, 'wb'))
+        data["estoi"].append(stoi(x, x_hat, sr_x, extended=True))
         data["si_sdr"].append(energy_ratios(x_hat, x, n)[0])
         data["si_sir"].append(energy_ratios(x_hat, x, n)[1])
         data["si_sar"].append(energy_ratios(x_hat, x, n)[2])