update

examples/clean_unet_aishell/run.sh

@@ -12,7 +12,7 @@ sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 

examples/clean_unet_aishell/step_2_train_model.py

@@ -32,7 +32,7 @@ from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanU
 from toolbox.torchaudio.models.clean_unet.modeling_clean_unet import CleanUNetPretrainedModel
 from toolbox.torchaudio.models.clean_unet.training import LinearWarmupCosineDecay
 from toolbox.torchaudio.models.clean_unet.loss import MultiResolutionSTFTLoss
-from toolbox.torchaudio.models.clean_unet.metrics import batch_pesq
+from toolbox.torchaudio.models.clean_unet.metrics import run_pesq_score
 
 torch.autograd.set_detect_anomaly(True)
 
@@ -217,7 +217,7 @@ def main():
         # train
         model.train()
 
-        total_pesq_metric = 0.
+        total_pesq_score = 0.
         total_loss = 0.
         total_ae_loss = 0.
         total_sc_loss = 0.
@@ -243,25 +243,21 @@ def main():
 
             enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
-            if pesq_metric is None:
-                pesq_metric = 0
-            else:
-                pesq_metric = torch.mean(pesq_metric).item()
+            pesq_score = run_pesq_score(enhanced_audios_list_r, clean_audios_list_r, sample_rate=8000, mode="nb")
 
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
             lr_scheduler.step()
 
-            total_pesq_metric += pesq_metric
+            total_pesq_score += pesq_score
             total_loss += loss.item()
             total_ae_loss += ae_loss.item()
             total_sc_loss += sc_loss.item()
             total_mag_loss += mag_loss.item()
             total_batches += 1
 
-            average_pesq_metric = round(total_pesq_metric / total_batches, 4)
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
             average_loss = round(total_loss / total_batches, 4)
             average_ae_loss = round(total_ae_loss / total_batches, 4)
             average_sc_loss = round(total_sc_loss / total_batches, 4)
@@ -269,7 +265,7 @@ def main():
 
             progress_bar.update(1)
             progress_bar.set_postfix({
-                "pesq_metric": average_pesq_metric,
+                "pesq_score": average_pesq_score,
                 "loss": average_loss,
                 "ae_loss": average_ae_loss,
                 "sc_loss": average_sc_loss,
@@ -281,7 +277,7 @@ def main():
 
         torch.cuda.empty_cache()
 
-        total_pesq_metric = 0.
+        total_pesq_score = 0.
         total_loss = 0.
         total_ae_loss = 0.
         total_sc_loss = 0.
@@ -308,20 +304,16 @@ def main():
 
                 enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
                 clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
-                if pesq_metric is None:
-                    pesq_metric = 0
-                else:
-                    pesq_metric = torch.mean(pesq_metric).item()
+                pesq_score = run_pesq_score(enhanced_audios_list_r, clean_audios_list_r, sample_rate=8000, mode="nb")
 
-                total_pesq_metric += pesq_metric
+                total_pesq_score += pesq_score
                 total_loss += loss.item()
                 total_ae_loss += ae_loss.item()
                 total_sc_loss += sc_loss.item()
                 total_mag_loss += mag_loss.item()
                 total_batches += 1
 
-                average_pesq_metric = round(total_pesq_metric / total_batches, 4)
+                average_pesq_score = round(total_pesq_score / total_batches, 4)
                 average_loss = round(total_loss / total_batches, 4)
                 average_ae_loss = round(total_ae_loss / total_batches, 4)
                 average_sc_loss = round(total_sc_loss / total_batches, 4)
@@ -329,7 +321,7 @@ def main():
 
                 progress_bar.update(1)
                 progress_bar.set_postfix({
-                    "pesq_metric": average_pesq_metric,
+                    "pesq_score": average_pesq_score,
                     "loss": average_loss,
                     "ae_loss": average_ae_loss,
                     "sc_loss": average_sc_loss,

toolbox/torchaudio/models/clean_unet/metrics.py

@@ -3,38 +3,76 @@
 from joblib import Parallel, delayed
 import numpy as np
 from pesq import pesq
-import torch
+from typing import List
 
+from pesq import cypesq
 
-def cal_pesq(clean, noisy, sr=16000):
+
+def run_pesq(clean_audio: np.ndarray,
+             noisy_audio: np.ndarray,
+             sample_rate: int = 16000,
+             mode: str = "wb",
+             ) -> float:
+    if sample_rate == 8000 and mode == "wb":
+        raise AssertionError(f"mode should be `nb` when sample_rate is 8000")
     try:
-        pesq_score = pesq(sr, clean, noisy, "wb")
+        pesq_score = pesq(sample_rate, clean_audio, noisy_audio, mode)
+    except cypesq.NoUtterancesError as e:
+        pesq_score = -1
     except Exception as e:
-        # print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
-        # error can happen due to silent period
+        print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
         pesq_score = -1
     return pesq_score
 
 
-def batch_pesq(clean, noisy):
-    pesq_score = Parallel(n_jobs=15)(delayed(cal_pesq)(c, n) for c, n in zip(clean, noisy))
-    pesq_score = np.array(pesq_score)
-    if -1 in pesq_score:
-        return None
-    pesq_score = (pesq_score - 1) / 3.5
-    return torch.FloatTensor(pesq_score)
+def run_batch_pesq(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+    parallel = Parallel(n_jobs=n_jobs)
+
+    parallel_tasks = list()
+    for clean_audio, noisy_audio in zip(clean_audio_list, noisy_audio_list):
+        parallel_task = delayed(run_pesq)(clean_audio, noisy_audio, sample_rate, mode)
+        parallel_tasks.append(parallel_task)
+
+    pesq_score_list = parallel.__call__(parallel_tasks)
+    return pesq_score_list
+
+
+def run_pesq_score(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+
+    pesq_score_list = run_batch_pesq(clean_audio_list=clean_audio_list,
+                                     noisy_audio_list=noisy_audio_list,
+                                     sample_rate=sample_rate,
+                                     mode=mode,
+                                     n_jobs=n_jobs,
+                                     )
+
+    pesq_score = np.mean(pesq_score_list)
+    return pesq_score
 
 
 def main():
+    clean_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+    noisy_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
 
-    prediction = torch.rand(size=(1, 160000), dtype=torch.float32)
-    ground_truth = torch.rand(size=(1, 160000), dtype=torch.float32)
+    clean_audio_list = list(clean_audio)
+    noisy_audio_list = list(noisy_audio)
 
-    prediction_list_r = list(prediction.cpu().numpy())
-    ground_truth_list_r = list(ground_truth.cpu().numpy())
+    pesq_score_list = run_batch_pesq(clean_audio_list, noisy_audio_list)
+    print(pesq_score_list)
 
-    pesq_score = batch_pesq(prediction_list_r, ground_truth_list_r)
+    pesq_score = run_pesq_score(clean_audio_list, noisy_audio_list)
     print(pesq_score)
+
     return
 
 

toolbox/torchaudio/models/mpnet/metrics.py

@@ -0,0 +1,80 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from joblib import Parallel, delayed
+import numpy as np
+from pesq import pesq
+from typing import List
+
+from pesq import cypesq
+
+
+def run_pesq(clean_audio: np.ndarray,
+             noisy_audio: np.ndarray,
+             sample_rate: int = 16000,
+             mode: str = "wb",
+             ) -> float:
+    if sample_rate == 8000 and mode == "wb":
+        raise AssertionError(f"mode should be `nb` when sample_rate is 8000")
+    try:
+        pesq_score = pesq(sample_rate, clean_audio, noisy_audio, mode)
+    except cypesq.NoUtterancesError as e:
+        pesq_score = -1
+    except Exception as e:
+        print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
+        pesq_score = -1
+    return pesq_score
+
+
+def run_batch_pesq(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+    parallel = Parallel(n_jobs=n_jobs)
+
+    parallel_tasks = list()
+    for clean_audio, noisy_audio in zip(clean_audio_list, noisy_audio_list):
+        parallel_task = delayed(run_pesq)(clean_audio, noisy_audio, sample_rate, mode)
+        parallel_tasks.append(parallel_task)
+
+    pesq_score_list = parallel.__call__(parallel_tasks)
+    return pesq_score_list
+
+
+def run_pesq_score(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+
+    pesq_score_list = run_batch_pesq(clean_audio_list=clean_audio_list,
+                                     noisy_audio_list=noisy_audio_list,
+                                     sample_rate=sample_rate,
+                                     mode=mode,
+                                     n_jobs=n_jobs,
+                                     )
+
+    pesq_score = np.mean(pesq_score_list)
+    return pesq_score
+
+
+def main():
+    clean_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+    noisy_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+
+    clean_audio_list = list(clean_audio)
+    noisy_audio_list = list(noisy_audio)
+
+    pesq_score_list = run_batch_pesq(clean_audio_list, noisy_audio_list)
+    print(pesq_score_list)
+
+    pesq_score = run_pesq_score(clean_audio_list, noisy_audio_list)
+    print(pesq_score)
+
+    return
+
+
+if __name__ == "__main__":
+    main()