Update app.py

app.py

@@ -6,69 +6,37 @@ from moviepy.editor import VideoFileClip, TextClip, CompositeVideoClip
 # Load the Whisper model
 model = whisper.load_model("base")  # Choose 'tiny', 'base', 'small', 'medium', or 'large'
 
-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)
-            f.write(f"{i + 1}\n")
-            f.write(f"{start_time} --> {end_time}\n")
-            f.write(f"{text.strip()}\n\n")
-
-def create_subtitled_video(video_file, srt_file):
-    # Load video
-    video_clip = VideoFileClip(video_file)
-
-    # Read SRT file and create text clips
-    subs = []
-    with open(srt_file, 'r') as f:
-        lines = f.readlines()
-        for i in range(0, len(lines), 4):
-            if i + 1 < len(lines):  # Ensure there is a next line for timestamps
-                timestamp_line = lines[i + 1].strip()
-                if " --> " in timestamp_line:  # Ensure it contains a timestamp
-                    try:
-                        start_time = convert_to_seconds(timestamp_line.split(" --> ")[0])
-                        end_time = convert_to_seconds(timestamp_line.split(" --> ")[1])
-                        text = lines[i + 2].strip() if (i + 2 < len(lines)) else ""
-                        text_clip = TextClip(text, fontsize=24, color='white', bg_color='black', size=video_clip.size, print_cmd=True)
-                        text_clip = text_clip.set_start(start_time).set_end(end_time).set_position('bottom')
-                        subs.append(text_clip)
-                    except Exception as e:
-                        print(f"Error processing subtitle {i // 4 + 1}: {e}")  # Debug print
-                        continue  # Skip to the next subtitle on error
-
-    # Combine video and subtitles
-    final_video = CompositeVideoClip([video_clip] + subs)
-    
-    # Save the final video with subtitles
-    subtitled_video_file = "subtitled_video.mp4"
-    final_video.write_videofile(subtitled_video_file, codec="libx264", audio_codec="aac")
-
-    return subtitled_video_file
-
-def convert_to_seconds(timestamp):
-    """Convert timestamp in SRT format (HH:MM:SS,mmm) to seconds."""
-    h, m, s = timestamp.split(':')
-    seconds = int(h) * 3600 + int(m) * 60 + float(s.replace(',', '.'))
-    return seconds
-
 def transcribe_video(video_file):
     # Transcribe the video to generate subtitles
     result = model.transcribe(video_file)
     
-    # Save the transcription to an .srt file
-    srt_file = "generated_subtitles.srt"
+    # 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)
     
-    # Write the transcription as subtitles
-    write_srt(result, srt_file)
+    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)
     
-    # Create the subtitled video
-    subtitled_video_file = create_subtitled_video(video_file, srt_file)
+    # Save the final video
+    subtitled_video_file = "subtitled_video.mp4"
+    final_video.write_videofile(subtitled_video_file, codec='libx264', audio_codec='aac')
     
     return subtitled_video_file
 
@@ -78,7 +46,7 @@ iface = gr.Interface(
     inputs=gr.File(label="Upload Video"),
     outputs=gr.File(label="Download Subtitled Video"),
     title="Video Subtitle Generator",
-    description="Upload a video file to generate subtitles and download a subtitled video."
+    description="Upload a video file to generate a subtitled video using Whisper."
 )
 
 if __name__ == "__main__":