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 1 --stop_stage 2 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 2 --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

@@ -243,7 +243,7 @@ def main():
 
             enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_score = run_pesq_score(enhanced_audios_list_r, clean_audios_list_r, sample_rate=8000, mode="nb")
+            pesq_score = run_pesq_score(clean_audios_list_r, enhanced_audios_list_r, sample_rate=8000, mode="nb")
 
             optimizer.zero_grad()
             loss.backward()
@@ -304,7 +304,7 @@ def main():
 
                 enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
                 clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                pesq_score = run_pesq_score(enhanced_audios_list_r, clean_audios_list_r, sample_rate=8000, mode="nb")
+                pesq_score = run_pesq_score(clean_audios_list_r, enhanced_audios_list_r, sample_rate=8000, mode="nb")
 
                 total_pesq_score += pesq_score
                 total_loss += loss.item()

examples/mpnet/step_2_train_model.py

@@ -26,9 +26,10 @@ from tqdm import tqdm
 
 from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
 from toolbox.torchaudio.models.mpnet.configuration_mpnet import MPNetConfig
-from toolbox.torchaudio.models.mpnet.discriminator import MetricDiscriminatorPretrainedModel, batch_pesq
-from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNet, MPNetPretrainedModel, phase_losses, pesq_score
+from toolbox.torchaudio.models.mpnet.discriminator import MetricDiscriminatorPretrainedModel
+from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNet, MPNetPretrainedModel, phase_losses
 from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft
+from toolbox.torchaudio.models.mpnet.metrics import run_batch_pesq, run_pesq_score
 
 
 def get_args():
@@ -251,7 +252,7 @@ def main():
             mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
 
             audio_list_r, audio_list_g = list(clean_audio.cpu().numpy()), list(audio_g.detach().cpu().numpy())
-            batch_pesq_score = batch_pesq(audio_list_r, audio_list_g)
+            pesq_score_list: List[float] = run_batch_pesq(audio_list_r, audio_list_g, sample_rate=config.sample_rate, mode="nb")
 
             # Discriminator
             optim_d.zero_grad()
@@ -259,11 +260,12 @@ def main():
             metric_g = discriminator.forward(clean_mag, mag_g_hat.detach())
             loss_disc_r = F.mse_loss(one_labels, metric_r.flatten())
 
-            if batch_pesq_score is not None:
-                loss_disc_g = F.mse_loss(batch_pesq_score.to(device), metric_g.flatten())
-            else:
-                # print("pesq is None!")
+            if -1 in pesq_score_list:
+                # print("-1 in batch_pesq_score!")
                 loss_disc_g = 0
+            else:
+                pesq_score_list: torch.FloatTensor = torch.tensor([(score - 1) / 3.5 for score in pesq_score_list], dtype=torch.float32)
+                loss_disc_g = F.mse_loss(pesq_score_list.to(device), metric_g.flatten())
 
             loss_disc_all = loss_disc_r + loss_disc_g
             loss_disc_all.backward()
@@ -334,11 +336,17 @@ def main():
                 audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
                 mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
 
-                total_pesq_score += pesq_score(
-                    torch.split(clean_audio, 1, dim=0),
-                    torch.split(audio_g, 1, dim=0),
-                    config
-                ).item()
+                clean_audio_list = torch.split(clean_audio, 1, dim=0)
+                enhanced_audio_list = torch.split(audio_g, 1, dim=0)
+                clean_audio_list = [t.squeeze().cpu().numpy() for t in clean_audio_list]
+                enhanced_audio_list = [t.squeeze().cpu().numpy() for t in enhanced_audio_list]
+                pesq_score = run_pesq_score(
+                    clean_audio_list,
+                    enhanced_audio_list,
+                    sample_rate = config.sample_rate,
+                    mode = "nb",
+                )
+                total_pesq_score += pesq_score
                 total_mag_err += F.mse_loss(clean_mag, mag_g).item()
                 val_ip_err, val_gd_err, val_iaf_err = phase_losses(clean_pha, pha_g)
                 total_pha_err += (val_ip_err + val_gd_err + val_iaf_err).item()