Update helper.py

helper.py

@@ -1,4 +1,5 @@
 import torch 
+import torchaudio.functional as F
 import matplotlib.pyplot as plt
 from pesq import pesq
 from pystoi import stoi
@@ -41,7 +42,19 @@ def si_snr(estimate, reference, epsilon=1e-8):
     si_snr = 10 * torch.log10(reference_pow) - 10 * torch.log10(error_pow)
     return si_snr.item()
 
+# def generate_mixture(waveform_clean, waveform_noise, target_snr):
+#     power_clean_signal = waveform_clean.pow(2).mean()
+#     power_noise_signal = waveform_noise.pow(2).mean()
+#     current_snr = 10 * torch.log10(power_clean_signal / power_noise_signal)
+#     waveform_noise *= 10 ** (-(target_snr - current_snr) / 20)
+#     return waveform_clean + waveform_noise
+
 def generate_mixture(waveform_clean, waveform_noise, target_snr):
+    if waveform_clean.size(1) > waveform_noise.size(1):
+        waveform_noise = F.pad(waveform_noise, (0, waveform_clean.size(1) - waveform_noise.size(1)))
+    elif waveform_noise.size(1) > waveform_clean.size(1):
+        waveform_clean = F.pad(waveform_clean, (0, waveform_noise.size(1) - waveform_clean.size(1)))
+
     power_clean_signal = waveform_clean.pow(2).mean()
     power_noise_signal = waveform_noise.pow(2).mean()
     current_snr = 10 * torch.log10(power_clean_signal / power_noise_signal)