Update app.py

app.py

@@ -1,70 +1,64 @@
-import gradio as gr 
-import os
-import sys
+import gradio as gr
 import subprocess
-#from moviepy.editor import VideoFileClip
-
+import os
 import whisper
 from whisper.utils import write_vtt
 
-model = whisper.load_model("medium")
-
-title = "Add Text/Caption to your YouTube Shorts - MultiLingual"
-
-def video2mp3(video_file, output_ext="mp3"):
-    filename, ext = os.path.splitext(video_file)
-    subprocess.call(["ffmpeg", "-y", "-i", video_file, f"{filename}.{output_ext}"], 
-                    stdout=subprocess.DEVNULL,
-                    stderr=subprocess.STDOUT)
-    return f"{filename}.{output_ext}"
-
-
-def translate(input_video):
-
-    audio_file = video2mp3(input_video)
+model = whisper.load_model('tiny')
+title = 'Add Captions(CC) to your videos'
+
+def convert_mp4_mp3(file, output="mp3"):
+    """
+    Convert the Input Video files to Audio files (MP4 -> MP3)
+    using FFMPEG 
+    """
+    filename, ext = os.path.splitext(file)
+    subprocess.call(['ffmpeg', '-y', '-i', file, f'{filename}.{output}'],
+                    stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
     
-    options = dict(beam_size=5, best_of=5, fp16 = False)
-    translate_options = dict(task="translate", **options)
-    result = model.transcribe(audio_file,**translate_options)
+    return f"{filename}.{output}"
+
+def transcribe(video):
+    """
+    Transcribe the text in the video file using Whisper model
+    and write the transcribed captions to the video
+    """
+    audio_file = convert_mp4_mp3(video)
+    # CFG
+    options = dict(beam_size=5, best_of=5, fp16=False)
+    translate_options = dict(task='translate', **options)
+    result = model.transcribe(audio_file, **translate_options)
 
     output_dir = ''
-    audio_path = audio_file.split(".")[0]
-
-    with open(os.path.join(output_dir, audio_path + ".vtt"), "w") as vtt:
-      write_vtt(result["segments"], file=vtt)
+    # audio_path = audio_file.split('.')[0]
+    audio_path = os.path.splitext(os.path.basename(audio_file))[0]
 
-    subtitle = audio_path + ".vtt"
-    output_video = audio_path + "_subtitled.mp4"
+    # Write Subtitle onto a .vtt file
+    with open(os.path.join(output_dir, audio_path + '.vtt'), 'w') as f:
+        write_vtt(result['segments'], file=f)
 
-    os.system(f"ffmpeg -i {input_video} -vf subtitles={subtitle} {output_video}")
+    # Write the subtitles on the input video
+    # subtitle = audio_path + '.vtt'
+    # output_video = audio_path + '_subtitled.mp4'
+    # os.system(f'ffmpeg -i {video} -vf subtitles={subtitle} {output_video}')
+    output_video = os.path.join(output_dir, f'{audio_path}_subtitled.mp4')
+    os.system(f'ffmpeg -i {video} -vf subtitles={os.path.join(output_dir, audio_path + ".vtt")} {output_video}')
 
     return output_video
 
 block = gr.Blocks()
 with block:
-
     with gr.Group():
         with gr.Box():
             with gr.Row().style():
-                inp_video = gr.Video(
+                input_video = gr.Video(
                     label="Input Video",
                     type="filepath",
-                    mirror_webcam = False
+                    mirror_webcam=False
                 )
-                op_video = gr.Video()
-            btn = gr.Button("Generate Subtitle Video")
-        
-        
-        
-
+                output_video = gr.Video()
+            btn = gr.Button('Generate Subtitle Video')
 
-        
-        btn.click(translate, inputs=[inp_video], outputs=[op_video])
-        gr.HTML('''
-        <div class="footer">
-                    <p>Powered by <a href="https://openai.com/" style="text-decoration: underline;" target="_blank">OpenAI</a> - Developed by <a href="https://github.com/Isaakkamau" style="text-decoration: underline;" target="_blank">+254704205553</a>
-                    </p>
-        </div>
-        ''')
+        btn.click(transcribe, inputs=[input_video], outputs=[output_video])
 
-block.launch(enable_queue = True)
+block.launch(enable_queue=True)