Upload noise_removal.py

utils/noise_removal.py

@@ -1,8 +1,13 @@
 from speechbrain.pretrained import SpectralMaskEnhancement
 import torchaudio
+import os
+import torch
 
-model = SpectralMaskEnhancement.from_hparams(source="speechbrain/metricgan-plus-voicebank")
-
+model = SpectralMaskEnhancement.from_hparams(
+    source="speechbrain/metricgan-plus-voicebank", savedir="tmp/denoiser"
+)
 def remove_noise(input_path, output_path):
-    enhanced = model.enhance_file(input_path)
-    torchaudio.save(output_path, enhanced[0], enhanced[1])
+    noisy_audio, sr = torchaudio.load(input_path)
+    lengths = torch.tensor([noisy_audio.shape[1]]) / noisy_audio.shape[1]
+    denoised = model.enhance_batch(noisy_audio, lengths)
+    torchaudio.save(output_path, denoised.cpu(), sr)