app.py

import gradio as gr
import whisper
import moviepy.editor as mp
from deep_translator import GoogleTranslator
from pydub import AudioSegment
import os
import tempfile

def extract_audio(video_path):
    video = mp.VideoFileClip(video_path)
    audio = video.audio
    audio_path = tempfile.mktemp(suffix=".wav")
    audio.write_audiofile(audio_path)
    return audio_path

def generate_subtitles(audio_path):
    model = whisper.load_model("base")
    result = model.transcribe(audio_path)
    return result["segments"]

def translate_subtitles(subtitles, target_language):
    translator = GoogleTranslator(source='auto', target=target_language)
    translated_subtitles = []
    for segment in subtitles:
        translated_text = translator.translate(segment["text"])
        translated_subtitles.append({
            "start": segment["start"],
            "end": segment["end"],
            "text": translated_text
        })
    return translated_subtitles

def add_subtitles_to_video(video_path, subtitles, output_path):
    video = mp.VideoFileClip(video_path)
    subtitles_clips = [
        mp.TextClip(txt=subtitle["text"], fontsize=24, color='white', bg_color='black', font='Arial')
        .set_position(('center', 'bottom'))
        .set_duration(subtitle["end"] - subtitle["start"])
        .set_start(subtitle["start"])
        for subtitle in subtitles
    ]
    final_video = mp.CompositeVideoClip([video] + subtitles_clips)
    final_video.write_videofile(output_path, codec="libx264", audio_codec="aac")

def process_video(video_path, target_language):
    # Extract audio from video
    audio_path = extract_audio(video_path)

    # Generate subtitles using Whisper
    subtitles = generate_subtitles(audio_path)

    # Translate subtitles
    translated_subtitles = translate_subtitles(subtitles, target_language)

    # Add translated subtitles to video
    output_path = tempfile.mktemp(suffix=".mp4")
    add_subtitles_to_video(video_path, translated_subtitles, output_path)

    return output_path

def gradio_interface(video, target_language):
    output_video = process_video(video.name, target_language)
    return output_video

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Video(label="Upload Video"),
        gr.Dropdown(choices=["es", "fr", "de", "it", "ja", "ko", "zh-CN"], label="Target Language")
    ],
    outputs=gr.Video(label="Processed Video"),
    title="Video Subtitle Translator",
    description="Upload a video, and get it back with translated subtitles!"
)

iface.launch()