+spectral gating filter

app.py

@@ -35,7 +35,12 @@ def preprocess_audio(audio_tensor, original_sample_rate, apply_normalization):
     audio_tensor = torchaudio.functional.resample(audio_tensor, orig_freq=original_sample_rate, new_freq=16000)  # Resample
     return audio_tensor
 
-def apply_wiener_filter(audio_tensor):
+def spectral_gating(audio_tensor):
+    audio_data = audio_tensor.numpy()
+    reduced_noise = nr.reduce_noise(y=audio_data, sr=16_000)
+    return torch.tensor(reduced_noise, dtype=audio_tensor.dtype)
+
+def wiener_filter(audio_tensor):
     audio_data = audio_tensor.numpy()
     filtered_audio = wiener(audio_data)
     return torch.tensor(filtered_audio, dtype=audio_tensor.dtype)
@@ -43,13 +48,13 @@ def apply_wiener_filter(audio_tensor):
 @spaces.GPU
 def transcribe_speech(audio, progress=gr.Progress()):
     if audio is None:
-        return "No audio received."
+        return "No audio received.", None
     progress(0.5, desc="Transcribing audio...")
     audio_np = audio.numpy().squeeze()
     transcription = pipe(audio_np, chunk_length_s=10)['text']
-    return replace_symbols_back(transcription)
+    return replace_symbols_back(transcription), audio
 
-def transcribe_from_youtube(url, apply_wiener, apply_normalization, progress=gr.Progress()):
+def transcribe_from_youtube(url, apply_wiener_filter, apply_normalization, apply_spectral_gating, progress=gr.Progress()):
     progress(0, "Downloading YouTube audio...")
     
     yt = YouTube(url)
@@ -62,16 +67,24 @@ def transcribe_from_youtube(url, apply_wiener, apply_normalization, progress=gr.
         audio, original_sample_rate = torchaudio.load(audio_data)
         audio = preprocess_audio(audio, original_sample_rate, apply_normalization)
 
-        if apply_wiener:
+        if apply_wiener_filter:
             progress(0.4, "Applying Wiener filter...")
-            audio = apply_wiener_filter(audio)
+            audio = wiener_filter(audio)
+
+        if apply_spectral_gating:
+            progress(0.4, "Applying Spectral Gating filter...")
+            audio = spectral_gating(audio)
 
-        transcription = transcribe_speech(audio)
+        transcription, processed_audio = transcribe_speech(audio)
+        audio_np = processed_audio.numpy().squeeze()
+        audio_output = BytesIO()
+        torchaudio.save(audio_output, torch.tensor(audio_np).unsqueeze(0), 16000)
+        audio_output.seek(0)
 
     except Exception as e:
-        return str(e)
+        return str(e), None
 
-    return transcription
+    return transcription, audio_output
 
 def populate_metadata(url):
     yt = YouTube(url)
@@ -96,16 +109,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         mic_audio = gr.Audio(sources=['microphone','upload'], type="filepath", label="Record or upload an audio")
         transcribe_button = gr.Button("Transcribe")
         transcription_output = gr.Textbox(label="Transcription")
+        audio_output = gr.Audio(label="Processed Audio")
 
-        transcribe_button.click(fn=transcribe_speech, inputs=mic_audio, outputs=transcription_output)
+        transcribe_button.click(fn=transcribe_speech, inputs=mic_audio, outputs=[transcription_output, audio_output])
 
     with gr.Tab("YouTube URL"):
         gr.Markdown("## Transcribe speech from YouTube video")
         youtube_url = gr.Textbox(label="Enter YouTube video URL")
 
         with gr.Accordion("Audio Improvements", open=False):
-            apply_wiener = gr.Checkbox(label="Reduce noise", info="Apply Wiener Filter", value=False)
             apply_normalization = gr.Checkbox(label="Normalize audio volume", value=True)
+            apply_spectral_gating = gr.Checkbox(label="Apply Spectral Gating filter", info="Noise reduction", value=True)
+            apply_wiener = gr.Checkbox(label="Apply Wiener filter", info="Noise reduction", value=False)
 
         with gr.Row():
             img = gr.Image(label="Thumbnail", height=240, width=240, scale=1)
@@ -113,8 +128,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
         transcribe_button = gr.Button("Transcribe")
         transcription_output = gr.Textbox(label="Transcription", placeholder="Transcription Output", lines=10)
+        audio_output = gr.Audio(label="Processed Audio")
 
-        transcribe_button.click(fn=transcribe_from_youtube, inputs=[youtube_url, apply_wiener, apply_normalization], outputs=transcription_output)
+        transcribe_button.click(fn=transcribe_from_youtube, inputs=[youtube_url, apply_wiener, apply_normalization, apply_spectral_gating], outputs=[transcription_output, audio_output])
         youtube_url.change(populate_metadata, inputs=[youtube_url], outputs=[img, title])
 
 demo.launch()