Update app.py

app.py

@@ -1,52 +1,46 @@
 import gradio as gr
 import whisper
 import os
-from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
 
 # Load the Whisper model
 model = whisper.load_model("base")  # Choose 'tiny', 'base', 'small', 'medium', or 'large'
 
-def transcribe_video(video_file):
+def write_srt(transcription, output_file):
+    with open(output_file, "w") as f:
+        for i, segment in enumerate(transcription['segments']):
+            start = segment['start']
+            end = segment['end']
+            text = segment['text']
+            # Format timestamps for SRT
+            start_time = whisper.utils.format_timestamp(start)
+            end_time = whisper.utils.format_timestamp(end)
+            print(f"Writing subtitle {i + 1}: {text.strip()} ({start_time} --> {end_time})")  # Debug print
+            f.write(f"{i + 1}\n")
+            f.write(f"{start_time} --> {end_time}\n")
+            f.write(f"{text.strip()}\n\n")
+
+def transcribe_video(video_file, language):
     # Transcribe the video to generate subtitles
-    result = model.transcribe(video_file)
-    
-    # Create a list of (start_time, end_time, text) tuples for subtitles
-    subtitles = [(segment['start'], segment['end'], segment['text'].strip()) for segment in result['segments']]
-    
-    # Create a subtitled video
-    subtitled_video_file = create_subtitled_video(video_file, subtitles)
+    result = model.transcribe(video_file, language=language)
     
-    return subtitled_video_file
-
-def create_subtitled_video(video_file, subtitles):
-    # Load the original video
-    video = VideoFileClip(video_file)
-
-    # Create a list of TextClips for each subtitle
-    text_clips = []
-    for start, end, text in subtitles:
-        text_clip = (TextClip(text, fontsize=24, color='white', bg_color='black', size=video.size)
-                     .set_start(start)
-                     .set_duration(end - start)
-                     .set_position(('center', 'bottom')))  # Position the subtitle at the bottom center
-        text_clips.append(text_clip)
-
-    # Overlay the subtitles on the video
-    final_video = CompositeVideoClip([video] + text_clips)
+    # Save the transcription to an .srt file
+    srt_file = "generated_subtitles.srt"
     
-    # Save the final video
-    subtitled_video_file = "subtitled_video.mp4"
-    final_video.write_videofile(subtitled_video_file, codec='libx264', audio_codec='aac')
+    # Write the transcription as subtitles
+    write_srt(result, srt_file)
     
-    return subtitled_video_file
+    return srt_file
 
 # Gradio interface
 iface = gr.Interface(
     fn=transcribe_video,
-    inputs=gr.File(label="Upload Video"),
-    outputs=gr.File(label="Download Subtitled Video"),
+    inputs=[
+        gr.File(label="Upload Video"),
+        gr.Dropdown(label="Select Language", choices=["en", "es", "fr", "de", "it", "pt"], value="en")
+    ],
+    outputs=gr.File(label="Download Subtitles"),
     title="Video Subtitle Generator",
-    description="Upload a video file to generate a subtitled video using Whisper."
+    description="Upload a video file to generate subtitles using Whisper. Select the language of the video."
 )
 
 if __name__ == "__main__":