r7 / app.py
RXTIME's picture
Update app.py
64e449c verified
raw
history blame
4.54 kB
import os
import subprocess
import soundfile as sf
from transformers import WhisperProcessor, WhisperForConditionalGeneration, MarianMTModel, MarianTokenizer
from gtts import gTTS
import gradio as gr
# Configurações
WHISPER_MODEL = "openai/whisper-medium"
TRANSLATION_MODEL = "Helsinki-NLP/opus-mt-en-pt"
TEMP_AUDIO_FILE = "temp_audio.wav"
TEMP_TRANSLATED_AUDIO = "translated_audio.mp3"
OUTPUT_VIDEO = "output_video.mp4"
def extract_audio(video_path: str, output_audio_path: str) -> None:
"""
Extrai o áudio de um vídeo usando FFmpeg.
"""
try:
command = [
"ffmpeg", "-i", video_path, "-q:a", "0", "-map", "a", output_audio_path, "-y"
]
subprocess.run(command, check=True)
except subprocess.CalledProcessError as e:
raise Exception(f"Erro ao extrair áudio: {e}")
def transcribe_audio(audio_path: str) -> str:
"""
Transcreve o áudio para texto usando o modelo Whisper.
"""
try:
processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL)
audio, _ = sf.read(audio_path)
input_features = processor(audio, sampling_rate=16000, return_tensors="pt").input_features
predicted_ids = model.generate(input_features)
transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
return transcription
except Exception as e:
raise Exception(f"Erro ao transcrever áudio: {e}")
def translate_text(text: str, source_lang: str = "en", target_lang: str = "pt") -> str:
"""
Traduz o texto para o idioma desejado usando o modelo MarianMT.
"""
try:
model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)
translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
translation = tokenizer.decode(translated[0], skip_special_tokens=True)
return translation
except Exception as e:
raise Exception(f"Erro ao traduzir texto: {e}")
def text_to_speech(text: str, output_audio_path: str, lang: str = "pt") -> None:
"""
Converte texto em áudio usando gTTS.
"""
try:
tts = gTTS(text, lang=lang)
tts.save(output_audio_path)
except Exception as e:
raise Exception(f"Erro ao gerar áudio: {e}")
def merge_audio_video(video_path: str, audio_path: str, output_path: str) -> None:
"""
Combina o áudio traduzido com o vídeo original usando FFmpeg.
"""
try:
command = [
"ffmpeg", "-i", video_path, "-i", audio_path, "-c:v", "copy", "-map", "0:v:0", "-map", "1:a:0", "-shortest", output_path, "-y"
]
subprocess.run(command, check=True)
except subprocess.CalledProcessError as e:
raise Exception(f"Erro ao combinar áudio e vídeo: {e}")
def translate_video(video_path: str) -> str:
"""
Função principal que orquestra a tradução do vídeo.
"""
try:
# 1. Extrair áudio do vídeo
extract_audio(video_path, TEMP_AUDIO_FILE)
# 2. Transcrever áudio para texto
transcription = transcribe_audio(TEMP_AUDIO_FILE)
print(f"Transcrição: {transcription}")
# 3. Traduzir texto para português
translation = translate_text(transcription)
print(f"Tradução: {translation}")
# 4. Converter texto traduzido em áudio
text_to_speech(translation, TEMP_TRANSLATED_AUDIO)
# 5. Combinar áudio traduzido com vídeo original
merge_audio_video(video_path, TEMP_TRANSLATED_AUDIO, OUTPUT_VIDEO)
# 6. Limpar arquivos temporários
os.remove(TEMP_AUDIO_FILE)
os.remove(TEMP_TRANSLATED_AUDIO)
return OUTPUT_VIDEO
except Exception as e:
print(f"Erro durante o processamento: {e}")
raise
# Interface Gradio
def gradio_interface(video):
try:
output_video = translate_video(video)
return output_video
except Exception as e:
return f"Erro: {e}"
# Configuração da interface
iface = gr.Interface(
fn=gradio_interface,
inputs=gr.Video(label="Upload do Vídeo"),
outputs=gr.Video(label="Vídeo Traduzido"),
title="Tradutor de Vídeos para Português",
description="Faça upload de um vídeo em qualquer idioma e receba o vídeo com áudio traduzido para português."
)
# Iniciar a interface
iface.launch()