Update app.py

app.py

@@ -1,20 +1,25 @@
 import gradio as gr
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-import numpy as np
-import pandas as pd
-import re
 from pydub import AudioSegment
 from pydub.generators import Sine
 import io
-# import torch
+import ffmpeg
+import subprocess
+import torch
+from moviepy.editor import VideoFileClip, AudioFileClip, CompositeVideoClip
+import tempfile
+import numpy as np
+import pandas as pd
+import re
+import scipy.io.wavfile
 
-# device = "cuda:0" if torch.cuda.is_available() else "cpu"
-model_id = "openai/whisper-large-v3"
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+model_id = "openai/whisper-tiny"
 
 model = AutoModelForSpeechSeq2Seq.from_pretrained(
     model_id, low_cpu_mem_usage=True, use_safetensors=True
 )
-# model.to(device)
+model.to(device)
 
 processor = AutoProcessor.from_pretrained(model_id)
 
@@ -34,6 +39,25 @@ arabic_bad_Words = pd.read_csv("arabic_bad_words_dataset.csv")
 english_bad_Words = pd.read_csv("english_bad_words_dataset.csv")
 
 
+def load_audio(file: str, sr: int = 16000):
+    try:
+        # This reads the audio from the video file without creating a separate audio file
+        command = [
+            "ffmpeg",
+            "-i", file,
+            "-f", "s16le",
+            "-acodec", "pcm_s16le",
+            "-ar", str(sr),
+            "-ac", "1",
+            "-"
+        ]
+        
+        out = subprocess.run(command, capture_output=True, check=True).stdout
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+    
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
 def clean_english_word(word):
     cleaned_text = re.sub(r'^[\s\W_]+|[\s\W_]+$', '', word)
     return cleaned_text.lower()
@@ -95,7 +119,7 @@ def format_output_to_list(data):
     formatted_list = "\n".join([f"{item['timestamp'][0]}s - {item['timestamp'][1]}s \t : {item['text']}" for item in data])
     return formatted_list
 
-def transcribe(input_audio, audio_language, task, timestamp_type):
+def transcribe_audio(input_audio, audio_language, task, timestamp_type):
     if input_audio is None:
         raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
 
@@ -129,31 +153,94 @@ def transcribe(input_audio, audio_language, task, timestamp_type):
 
     return  [text, timestamps, foul_words, (sample_rate, audio_data)]
 
-cache_examples = [
-    ["arabic_english_audios/audios/english_audio_18.mp3", 'English', 'transcribe', 'word'],
-    ["arabic_english_audios/audios/english_audio_20.mp3", 'English', 'transcribe', 'word'],
-    ["arabic_english_audios/audios/english_audio_21.mp3", 'English', 'transcribe', 'word'],
-    ["arabic_english_audios/audios/english_audio_22.mp3", 'English', 'transcribe', 'word'],
-    ["arabic_english_audios/audios/english_audio_27.mp3", 'English', 'transcribe', 'word'],
-    ["arabic_english_audios/audios/english_audio_29.mp3", 'English', 'transcribe', 'word'],
-]
+
+def transcribe_video(input_video, video_language, task, timestamp_type):
+    # Load the video file
+    video = VideoFileClip(input_video)
+
+    # Extract the audio
+    audio = video.audio
+
+    # Create a temporary file to save the audio
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file:
+        audio.write_audiofile(temp_audio_file.name, codec='pcm_s16le')
+
+    # Load the audio file into an AudioSegment
+    audio_segment = AudioSegment.from_file(temp_audio_file.name, format="wav")
+
+    # Ensure the audio is mono
+    if audio_segment.channels > 1:
+        audio_segment = audio_segment.set_channels(1)
+
+    # Save the mono audio to a temporary file
+    mono_temp_audio_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    audio_segment.export(mono_temp_audio_file.name, format="wav")
+
+    # Save the mono audio to a file
+    extracted_audio_path = "extracted_audio_mono.mp3"
+    audio_segment.export(extracted_audio_path, format="mp3")
+
+    # Load the audio as a numpy array
+    # sample_rate, audio_array = scipy.io.wavfile.read(mono_temp_audio_file.name)
+
+    output = pipe(extracted_audio_path, return_timestamps=timestamp_type, generate_kwargs={"task": task})
+    text = output['text']
+
+    timestamps = format_output_to_list(output['chunks'])
+
+    foul_words, negative_timestamps = classifier(output['chunks'], video_language)
+    foul_words = ", ".join(foul_words)
+
+
+    audio_output = mute_audio_range(extracted_audio_path, negative_timestamps)
+
+    # Resample the output audio to 16kHz
+    audio_output = resample_audio(audio_output, 16000)
+
+    # Save the output audio to a BytesIO object
+    output_buffer = io.BytesIO()
+    audio_output.export(output_buffer, format="wav")
+    output_buffer.seek(0)
+    
+    # Read the audio data from the BytesIO buffer
+    sample_rate = audio_output.frame_rate
+    audio_data = np.frombuffer(output_buffer.read(), dtype=np.int16)
+
+    # Save the processed NumPy array to a WAV file
+    processed_audio_path = "processed_audio.wav"
+    scipy.io.wavfile.write(processed_audio_path, sample_rate, audio_data)
+
+    # Load the processed audio into a moviepy AudioFileClip
+    processed_audio = AudioFileClip(processed_audio_path)
+
+    # Set the audio of the video to the processed audio
+    final_video = video.set_audio(processed_audio)
+
+    # Save the final video with the combined audio
+    final_video_path = "final_video_with_processed_audio.mp4"
+    final_video.write_videofile(final_video_path, codec="libx264", audio_codec="aac")
+
+
+    return [text, timestamps, foul_words, final_video_path]
+
+
 examples = [
         ["arabic_english_audios/audios/arabic_audio_11.mp3", 'Arabic', 'transcribe', 'word'],
         ["arabic_english_audios/audios/arabic_audio_12.mp3", 'Arabic', 'transcribe', 'word'],
         ["arabic_english_audios/audios/arabic_audio_13.mp3", 'Arabic', 'transcribe', 'word'],
     
-        # ["arabic_english_audios/audios/english_audio_18.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_18.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_19.mp3", 'English', 'transcribe', 'word'],
-        # ["arabic_english_audios/audios/english_audio_20.mp3", 'English', 'transcribe', 'word'],
-        # ["arabic_english_audios/audios/english_audio_21.mp3", 'English', 'transcribe', 'word'],
-        # ["arabic_english_audios/audios/english_audio_22.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_20.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_21.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_22.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_23.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_24.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_25.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_26.mp3", 'English', 'transcribe', 'word'],
-        # ["arabic_english_audios/audios/english_audio_27.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_27.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_28.mp3", 'English', 'transcribe', 'word'],
-        # ["arabic_english_audios/audios/english_audio_29.mp3", 'English', 'transcribe', 'word'],
+        ["arabic_english_audios/audios/english_audio_29.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_30.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_31.mp3", 'English', 'transcribe', 'word'],
         ["arabic_english_audios/audios/english_audio_32.mp3", 'English', 'transcribe', 'word'],
@@ -175,27 +262,51 @@ examples = [
 with gr.Blocks(theme=gr.themes.Default()) as demo:
     gr.HTML("<h2 style='text-align: center;'>Transcribing Audio with Timestamps using whisper-large-v3</h2>")
     # gr.Markdown("")
-    with gr.Row():
-        with gr.Column():
-            audio_input = gr.Audio(sources=["upload", 'microphone'], type="filepath", label="Audio file")
-            audio_language = gr.Radio(["Arabic", "English"], label="Audio Language")
-            task = gr.Radio(["transcribe", "translate"], label="Task")
-            timestamp_type = gr.Radio(["sentence", "word"], label="Timestamp Type")
-            with gr.Row():
-                clear_button = gr.ClearButton(value="Clear")
-                submit_button = gr.Button("Submit", variant="primary", )
-
-        with gr.Column():
-            transcript_output = gr.Text(label="Transcript")
-            timestamp_output = gr.Text(label="Timestamps")
-            foul_words = gr.Text(label="Foul Words")
-            output_audio = gr.Audio(label="Output Audio", type="numpy")
-            
-    cache_examples = gr.Examples(cache_examples, inputs=[audio_input, audio_language, task, timestamp_type], outputs=[transcript_output, timestamp_output,  foul_words, output_audio], fn=transcribe, examples_per_page=10,  cache_examples=True, label="Sample Examples")
-    non_cache_examples = gr.Examples(examples, inputs=[audio_input, audio_language, task, timestamp_type], outputs=[transcript_output, timestamp_output,  foul_words, output_audio], fn=transcribe, examples_per_page=50,  cache_examples=False, label="Other Examples")
-
-    submit_button.click(fn=transcribe, inputs=[audio_input, audio_language, task, timestamp_type], outputs=[transcript_output, timestamp_output, foul_words, output_audio])
-    clear_button.add([audio_input, audio_language, task, timestamp_type, transcript_output, timestamp_output, foul_words, output_audio])
+    with gr.Tab("Audio"):
+        with gr.Row():
+            with gr.Column():
+                audio_input = gr.Audio(sources=["upload", 'microphone'], type="filepath", label="Audio file")
+                audio_language = gr.Radio(["Arabic", "English"], label="Audio Language")
+                audio_task = gr.Radio(["transcribe", "translate"], label="Task")
+                audio_timestamp_type = gr.Radio(["sentence", "word"], label="Timestamp Type")
+                with gr.Row():
+                    audio_clear_button = gr.ClearButton(value="Clear")
+                    audio_submit_button = gr.Button("Submit", variant="primary", )
+
+            with gr.Column():
+                audio_transcript_output = gr.Text(label="Transcript")
+                audio_timestamp_output = gr.Text(label="Timestamps")
+                audio_foul_words = gr.Text(label="Foul Words")
+                output_audio = gr.Audio(label="Output Audio", type="numpy")
+
+        examples = gr.Examples(examples, inputs=[audio_input, audio_language, audio_task, audio_timestamp_type], outputs=[audio_transcript_output, audio_timestamp_output,  audio_foul_words, output_audio], fn=transcribe_audio, examples_per_page=50,  cache_examples=False)
+
+        audio_submit_button.click(fn=transcribe_audio, inputs=[audio_input, audio_language, audio_task, audio_timestamp_type], outputs=[audio_transcript_output, audio_timestamp_output, audio_foul_words, output_audio])
+        audio_clear_button.add([audio_input, audio_language, audio_task, audio_timestamp_type, audio_transcript_output, audio_timestamp_output, audio_foul_words, output_audio])
+
+
+    with gr.Tab("Video"):
+        with gr.Row():
+            with gr.Column():
+                video_input = gr.Video(sources=["upload", 'webcam'], label="Video file")
+                video_language = gr.Radio(["Arabic", "English"], label="Video Language")
+                video_task = gr.Radio(["transcribe", "translate"], label="Task")
+                video_timestamp_type = gr.Radio(["sentence", "word"], label="Timestamp Type")
+                with gr.Row():
+                    video_clear_button = gr.ClearButton(value="Clear")
+                    video_submit_button = gr.Button("Submit", variant="primary", )
+
+            with gr.Column():
+                video_transcript_output = gr.Text(label="Transcript")
+                video_timestamp_output = gr.Text(label="Timestamps")
+                video_foul_words = gr.Text(label="Foul Words")
+                output_video = gr.Video(label="Output Video")
+                # output_video = gr.Audio(label="Output Audio", type="numpy")
+
+
+        video_submit_button.click(fn=transcribe_video, inputs=[video_input, video_language, video_task, video_timestamp_type], outputs=[video_transcript_output, video_timestamp_output, video_foul_words, output_video])
+        video_clear_button.add([video_input, video_language, video_task, video_timestamp_type, video_transcript_output, video_timestamp_output, video_foul_words, output_video])
+
 
 
 if __name__ == "__main__":