update

examples/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -111,7 +111,6 @@ def enhance(mix_spec_complex: torch.Tensor,
     # print(f"speech_spec_prediction: {speech_spec_prediction.shape}")
     # print(f"noise_spec: {noise_spec.shape}")
 
-    speech_spec_prediction = torch.view_as_complex(speech_spec_prediction)
     speech_wave = istft.forward(speech_spec_prediction)
     # speech_wave = istft.forward(speech_spec)
     noise_wave = istft.forward(noise_spec)
@@ -245,6 +244,8 @@ def main():
 
         with torch.no_grad():
             speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_spec_complex)
+            speech_spec_prediction = torch.view_as_complex(speech_spec_prediction)
+
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
             # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
             # loss = irm_loss + 0.1 * snr_loss
@@ -252,14 +253,31 @@ def main():
 
         # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
         # speech_irm_prediction shape: [batch_size, freq_dim (256), time_steps]
-        # batch_size, _, time_steps = speech_irm_prediction.shape
-        # speech_irm_prediction = torch.concat(
-        #     [
-        #         speech_irm_prediction,
-        #         0.5*torch.ones(size=(batch_size, 1, time_steps), dtype=speech_irm_prediction.dtype).to(device)
-        #     ],
-        #     dim=1,
-        # )
+        batch_size, _, time_steps = speech_irm_prediction.shape
+
+
+        mix_spec_complex = torch.concat(
+            [
+                mix_spec_complex,
+                torch.zeros(size=(batch_size, 1, time_steps), dtype=mix_spec_complex.dtype).to(device)
+            ],
+            dim=1,
+        )
+        speech_spec_prediction = torch.concat(
+            [
+                speech_spec_prediction,
+                torch.zeros(size=(batch_size, 1, time_steps), dtype=speech_spec_prediction.dtype).to(device)
+            ],
+            dim=1,
+        )
+        speech_irm_prediction = torch.concat(
+            [
+                speech_irm_prediction,
+                0.5 * torch.ones(size=(batch_size, 1, time_steps), dtype=speech_irm_prediction.dtype).to(device)
+            ],
+            dim=1,
+        )
+
         # speech_irm_prediction shape: [batch_size, freq_dim (257), time_steps]
         speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_spec_prediction, speech_irm_prediction)
         save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)