Spaces:
Sleeping
Sleeping
File size: 3,390 Bytes
4dbcfe1 f475b61 4dbcfe1 2986f68 4dbcfe1 2986f68 4dbcfe1 f475b61 4dbcfe1 f475b61 4dbcfe1 f475b61 4dbcfe1 f475b61 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import gradio as gr
import moviepy.editor as mp
from deep_translator import GoogleTranslator
from pydub import AudioSegment
import os
import tempfile
import torch
from transformers import WhisperProcessor, WhisperForConditionalGeneration
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):
device = "cuda" if torch.cuda.is_available() else "cpu"
processor = WhisperProcessor.from_pretrained("openai/whisper-base")
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").to(device)
# Load and preprocess the audio
audio_input, _ = librosa.load(audio_path, sr=16000)
input_features = processor(audio_input, sampling_rate=16000, return_tensors="pt").input_features.to(device)
# Generate token ids
predicted_ids = model.generate(input_features)
# Decode token ids to text
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
# For simplicity, we're returning a single segment with the full transcription
# In a more advanced implementation, you might want to split this into multiple segments
return [{"start": 0, "end": len(audio_input) / 16000, "text": transcription[0]}]
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() |