Add audio quality improvements for youtube audios.

app.py

@@ -7,6 +7,10 @@ from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
 from pytube import YouTube
 from transformers import pipeline
 import re
+from pydub import AudioSegment
+from scipy.io import wavfile
+from scipy.signal import wiener
+import numpy as np
 
 # pipe = pipeline(model="anzorq/w2v-bert-2.0-kbd", device=0) # old model
 pipe = pipeline(model="anzorq/w2v-bert-2.0-kbd-v2", device=0) # new model with a new tokenizer
@@ -24,25 +28,50 @@ reverse_pattern = re.compile('|'.join(re.escape(key) for key in reverse_replacem
 def replace_symbols_back(text):
     return reverse_pattern.sub(lambda match: reverse_replacements[match.group(0)], text)
 
+def normalize_audio(audio_path):
+    audio = AudioSegment.from_file(audio_path, format="mp4")
+    normalized_audio = audio.normalize()
+    normalized_audio.export(audio_path, format="mp4")
+
+def apply_wiener_filter(audio_path):
+    sample_rate, audio_data = wavfile.read(audio_path)
+    filtered_audio = wiener(audio_data)
+    wavfile.write(audio_path, sample_rate, filtered_audio.astype(np.int16))
+
+def resample_audio(audio_path, target_sample_rate=16000):
+    audio, sample_rate = torchaudio.load(audio_path)
+    resampled_audio = torchaudio.transforms.Resample(sample_rate, target_sample_rate)(audio)
+    torchaudio.save(audio_path, resampled_audio, target_sample_rate)
+
 @spaces.GPU
 def transcribe_speech(audio, progress=gr.Progress()):
     if audio is None:  # Handle the NoneType error for microphone input
         return "No audio received."
-    
+
     progress(0.5, desc="Transcribing audio...")
     transcription = pipe(audio, chunk_length_s=10)['text']
-    
+
     return replace_symbols_back(transcription)
 
-def transcribe_from_youtube(url, progress=gr.Progress()):
+def transcribe_from_youtube(url, apply_improvements, progress=gr.Progress()):
     progress(0, "Downloading YouTube audio...")
-    # Download audio from YouTube using pytube
     audio_path = YouTube(url).streams.filter(only_audio=True)[0].download(filename="tmp.mp4")
 
+    if apply_improvements:
+        progress(0.2, "Normalizing audio...")
+        normalize_audio(audio_path)
+
+        progress(0.4, "Applying Wiener filter...")
+        apply_wiener_filter(audio_path)
+
+        progress(0.6, "Resampling audio...")
+        resample_audio(audio_path)
+
+    progress(0.8, "Transcribing audio...")
     transcription = transcribe_speech(audio_path)
 
     os.remove(audio_path)
-    
+
     return transcription
 
 def populate_metadata(url):
@@ -62,27 +91,28 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             </div>
         """
     )
-    
+
     with gr.Tab("Microphone Input"):
         gr.Markdown("## Transcribe speech from microphone")
         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")
-        
+
         transcribe_button.click(fn=transcribe_speech, inputs=mic_audio, outputs=transcription_output)
 
     with gr.Tab("YouTube URL"):
         gr.Markdown("## Transcribe speech from YouTube video")
         youtube_url = gr.Textbox(label="Enter YouTube video URL")
-        
+        apply_improvements = gr.Checkbox(label="Apply Audio Improvements", value=True)
+
         with gr.Row():
             img = gr.Image(label="Thumbnail", height=240, width=240, scale=1)
             title = gr.Label(label="Video Title", scale=2)
-        
+
         transcribe_button = gr.Button("Transcribe")
         transcription_output = gr.Textbox(label="Transcription", placeholder="Transcription Output", lines=10)
-        
-        transcribe_button.click(fn=transcribe_from_youtube, inputs=youtube_url, outputs=transcription_output)
+
+        transcribe_button.click(fn=transcribe_from_youtube, inputs=[youtube_url, apply_improvements], outputs=transcription_output)
         youtube_url.change(populate_metadata, inputs=[youtube_url], outputs=[img, title])
 
 demo.launch()