Files
criar_legendas/main.py
2024-11-26 10:20:07 -08:00

458 lines
17 KiB
Python

import tkinter as tk
from tkinter import ttk, scrolledtext, filedialog, messagebox
import threading
import whisper
import cv2
from datetime import timedelta
import os
import subprocess
import numpy as np
from transformers import pipeline
import torch
NLLB_LANG_CODES = {
'pt': 'por_Latn',
'en': 'eng_Latn',
'es': 'spa_Latn',
'fr': 'fra_Latn',
'de': 'deu_Latn',
'it': 'ita_Latn',
'nl': 'nld_Latn',
'pl': 'pol_Latn',
'ru': 'rus_Cyrl',
'zh': 'zho_Hans'
}
def check_dependencies():
try:
import torch
import transformers
import whisper
return True
except ImportError as e:
messagebox.showerror("Erro de Dependência",
"Por favor, instale todas as dependências necessárias:\n"
"pip install torch transformers whisper")
return False
class VideoSubtitleApp:
def __init__(self, root):
self.root = root
self.root.title("Extrator e Tradutor de Legendas")
self.root.geometry("900x700")
# Variáveis
self.video_path = tk.StringVar()
self.video_info = tk.StringVar()
self.selected_language = tk.StringVar(value='English')
self.translation_language = tk.StringVar(value='Português')
self.status_var = tk.StringVar(value="Pronto")
self.subtitles_list = []
# Dicionário de línguas para extração
self.languages = {
'Português (Brasil)': 'pt',
'Português (Portugal)': 'pt',
'English': 'en',
'Español': 'es',
'Français': 'fr',
'Deutsch': 'de',
'Italiano': 'it'
}
# Dicionário para tradução
self.translation_languages = {
'Português': 'pt',
'English': 'en',
'Español': 'es',
'Français': 'fr',
'Deutsch': 'de',
'Italiano': 'it',
'Nederlands': 'nl',
'Polski': 'pl',
'Русский': 'ru',
'中文': 'zh'
}
# Criar interface
self.create_widgets()
# Inicializar modelos em threads separadas
self.model_ready = False
self.translator = None
threading.Thread(target=self.initialize_models, daemon=True).start()
def initialize_models(self):
"""Inicializa os modelos de Whisper e Tradução"""
try:
self.status_var.set("Carregando modelos...")
# Inicializar Whisper
self.model = whisper.load_model("base")
# Usar modelo NLLB mais leve
device = 0 if torch.cuda.is_available() else -1
self.translator = pipeline(
"translation",
model="facebook/nllb-200-distilled-600M",
device=device
)
self.model_ready = True
self.status_var.set("Modelos carregados com sucesso")
self.generate_button.config(state='normal')
except Exception as e:
self.status_var.set("Erro ao carregar modelos")
messagebox.showerror("Erro", f"Erro ao carregar modelos:\n{str(e)}")
def create_widgets(self):
"""Cria a interface gráfica"""
# Frame principal
main_frame = ttk.Frame(self.root, padding=(10, 10, 10, 10))
main_frame.grid(row=0, column=0, sticky="nsew")
# Configurar grid
self.root.grid_rowconfigure(0, weight=1)
self.root.grid_columnconfigure(0, weight=1)
main_frame.grid_columnconfigure(0, weight=1)
# Frame superior
top_frame = ttk.Frame(main_frame)
top_frame.grid(row=0, column=0, sticky="ew", pady=(0, 10))
# Botões e controles
ttk.Button(top_frame, text="📂 Selecionar Vídeo",
command=self.select_file).pack(side=tk.LEFT, padx=5)
ttk.Label(top_frame, text="🌐 Idioma Original:").pack(side=tk.LEFT, padx=5)
language_combo = ttk.Combobox(top_frame,
values=list(self.languages.keys()),
textvariable=self.selected_language,
state='readonly',
width=20)
language_combo.pack(side=tk.LEFT, padx=5)
ttk.Label(top_frame, text="🔄 Traduzir para:").pack(side=tk.LEFT, padx=5)
translation_combo = ttk.Combobox(top_frame,
values=list(self.translation_languages.keys()),
textvariable=self.translation_language,
state='readonly',
width=20)
translation_combo.pack(side=tk.LEFT, padx=5)
# Caminho do arquivo
path_frame = ttk.LabelFrame(main_frame, text="Arquivo Selecionado")
path_frame.grid(row=1, column=0, sticky="ew", pady=(0, 10))
ttk.Label(path_frame, textvariable=self.video_path,
wraplength=800).grid(row=0, column=0, padx=5, pady=5)
# Informações do vídeo
info_frame = ttk.LabelFrame(main_frame, text="Informações do Vídeo")
info_frame.grid(row=2, column=0, sticky="ew", pady=(0, 10))
ttk.Label(info_frame, textvariable=self.video_info).grid(row=0, column=0, padx=5, pady=5)
# Botões de ação
button_frame = ttk.Frame(main_frame)
button_frame.grid(row=3, column=0, pady=(0, 10))
self.generate_button = ttk.Button(button_frame, text="🎬 Gerar Legendas",
command=self.generate_subtitles,
state='disabled')
self.generate_button.pack(side=tk.LEFT, padx=5)
self.translate_button = ttk.Button(button_frame, text="🔄 Traduzir",
command=self.translate_subtitles,
state='disabled')
self.translate_button.pack(side=tk.LEFT, padx=5)
self.save_button = ttk.Button(button_frame, text="💾 Salvar",
command=self.save_subtitles,
state='disabled')
self.save_button.pack(side=tk.LEFT, padx=5)
# Barra de progresso
self.progress = ttk.Progressbar(main_frame, mode='indeterminate')
self.progress.grid(row=4, column=0, sticky="ew", pady=(0, 10))
# Editor de legendas
editor_frame = ttk.LabelFrame(main_frame, text="Editor de Legendas")
editor_frame.grid(row=5, column=0, sticky="nsew", pady=(0, 10))
editor_frame.grid_columnconfigure(0, weight=1)
editor_frame.grid_rowconfigure(0, weight=1)
self.subtitle_text = scrolledtext.ScrolledText(
editor_frame,
wrap=tk.WORD,
font=('Consolas', 10)
)
self.subtitle_text.grid(row=0, column=0, sticky="nsew", padx=5, pady=5)
# Barra de status
status_frame = ttk.Frame(main_frame)
status_frame.grid(row=6, column=0, sticky="ew")
ttk.Label(status_frame, textvariable=self.status_var, relief=tk.SUNKEN).grid(
row=0, column=0, sticky="ew")
status_frame.grid_columnconfigure(0, weight=1)
def translate_subtitles(self):
"""Traduz as legendas para o idioma selecionado"""
if not self.subtitles_list:
messagebox.showwarning("Aviso", "Gere as legendas primeiro antes de traduzir.")
return
target_lang = self.translation_languages[self.translation_language.get()]
if self.selected_language.get() == self.translation_language.get():
messagebox.showwarning("Aviso", "O idioma de origem e destino são os mesmos.")
return
self.progress.start()
self.translate_button.config(state='disabled')
threading.Thread(target=self._translate_process, daemon=True).start()
def _translate_process(self):
"""Processo de tradução em background"""
try:
self.status_var.set("Traduzindo legendas...")
# Obter texto atual
current_text = self.subtitle_text.get('1.0', tk.END).strip()
if not current_text:
raise Exception("Nenhum texto para traduzir")
# Obter idioma alvo
target_lang = self.translation_languages[self.translation_language.get()]
target_lang_code = NLLB_LANG_CODES.get(target_lang, 'por_Latn') # Português como fallback
# Dividir em segmentos para preservar formato SRT
segments = current_text.split('\n\n')
translated_segments = []
total_segments = len(segments)
for i, segment in enumerate(segments, 1):
lines = segment.split('\n')
if len(lines) >= 3: # Verifica se é um segmento válido
# Traduz apenas o texto, mantém número e timestamp
text_to_translate = '\n'.join(lines[2:])
try:
translation = self.translator(
text_to_translate,
src_lang="eng_Latn",
tgt_lang=target_lang_code,
max_length=512
)[0]['translation_text']
print(f"Original: {text_to_translate}")
print(f"Tradução: {translation}")
except Exception as e:
print(f"Erro ao traduzir segmento {i}: {str(e)}")
translation = text_to_translate # mantém texto original em caso de erro
# Reconstrói o segmento
translated_segment = f"{lines[0]}\n{lines[1]}\n{translation}"
translated_segments.append(translated_segment)
# Atualiza status
self.status_var.set(f"Traduzindo... {i}/{total_segments}")
# Atualiza o texto na interface
translated_text = '\n\n'.join(translated_segments)
self.root.after(0, self._update_translation, translated_text)
self.status_var.set("Tradução concluída!")
except Exception as e:
self.status_var.set("Erro na tradução")
messagebox.showerror("Erro", f"Erro ao traduzir legendas: {str(e)}")
finally:
self.progress.stop()
self.translate_button.config(state='normal')
# [Resto dos métodos existentes permanece igual...]
def select_file(self):
"""Seleciona arquivo de vídeo"""
filename = filedialog.askopenfilename(
title="Selecionar Vídeo",
filetypes=[
("Arquivos de Vídeo", "*.mp4 *.mkv *.avi"),
("Todos os Arquivos", "*.*")
]
)
if filename:
self.video_path.set(filename)
self.load_video_info(filename)
def load_video_info(self, filename):
"""Carrega informações do vídeo"""
try:
cap = cv2.VideoCapture(filename)
fps = cap.get(cv2.CAP_PROP_FPS)
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
duration = frame_count / fps
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
info = f"Duração: {str(timedelta(seconds=int(duration)))}\n"
info += f"Resolução: {width}x{height}\n"
info += f"FPS: {fps:.2f}\n"
info += f"Formato: {os.path.splitext(filename)[1]}"
self.video_info.set(info)
cap.release()
except Exception as e:
messagebox.showerror("Erro", f"Erro ao carregar vídeo: {str(e)}")
def extract_audio(self, video_path, audio_path):
"""Extrai o áudio do vídeo"""
try:
command = [
'ffmpeg',
'-i', video_path,
'-vn',
'-acodec', 'pcm_s16le',
'-ar', '16000',
'-ac', '1',
'-y',
audio_path
]
process = subprocess.run(
command,
capture_output=True,
text=True
)
return os.path.exists(audio_path) and os.path.getsize(audio_path) > 0
except Exception as e:
print(f"Erro na extração de áudio: {str(e)}")
return False
def process_audio_with_whisper(self, audio_path, language_code):
"""Processa o áudio usando Whisper"""
try:
# Configurar opções do Whisper
options = {
"language": language_code,
"task": "transcribe",
"verbose": False
}
# Realizar transcrição
result = self.model.transcribe(audio_path, **options)
# Processar segmentos
segments = []
for i, segment in enumerate(result["segments"], 1):
start_time = segment["start"]
end_time = segment["end"]
text = segment["text"].strip()
if text:
segment_str = f"{i}\n"
segment_str += f"{self.format_timestamp(start_time)} --> {self.format_timestamp(end_time)}\n"
segment_str += f"{text}\n\n"
segments.append(segment_str)
return segments
except Exception as e:
print(f"Erro detalhado no processamento do áudio: {str(e)}")
raise
def format_timestamp(self, seconds):
"""Converte segundos para formato SRT"""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
millisecs = int((seconds * 1000) % 1000)
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millisecs:03d}"
def generate_subtitles(self):
"""Inicia processo de geração de legendas"""
if not self.video_path.get():
messagebox.showwarning("Aviso", "Selecione um vídeo primeiro.")
return
if not self.model_ready:
messagebox.showwarning("Aviso", "Aguarde o modelo ser carregado.")
return
self.progress.start()
self.generate_button.config(state='disabled')
self.save_button.config(state='disabled')
threading.Thread(target=self.process_video, daemon=True).start()
def process_video(self):
"""Processa o vídeo e gera legendas"""
audio_path = "temp_audio.wav"
try:
self.status_var.set("Extraindo áudio...")
if not self.extract_audio(self.video_path.get(), audio_path):
raise Exception("Falha na extração do áudio")
self.status_var.set("Processando áudio...")
language = self.languages[self.selected_language.get()]
self.subtitles_list = self.process_audio_with_whisper(audio_path, language)
if not self.subtitles_list:
raise Exception("Nenhuma legenda gerada")
self.status_var.set("Legendas geradas com sucesso!")
self.root.after(0, self.update_subtitle_text)
except Exception as e:
self.status_var.set("Erro no processamento")
messagebox.showerror("Erro", str(e))
finally:
self.progress.stop()
self.generate_button.config(state='normal')
self.translate_button.config(state='normal')
try:
if os.path.exists(audio_path):
os.remove(audio_path)
except:
pass
def update_subtitle_text(self):
"""Atualiza o texto das legendas na interface"""
self.subtitle_text.delete('1.0', tk.END)
self.subtitle_text.insert('1.0', ''.join(self.subtitles_list))
self.save_button.config(state='normal')
self.translate_button.config(state='normal')
def save_subtitles(self):
"""Salva as legendas em arquivo"""
try:
# Determina sufixo baseado no idioma
language_code = self.translation_languages.get(self.translation_language.get(), 'en')
suffix = f"_{language_code}" if self.translation_language.get() != 'English' else ""
# Gera nome do arquivo
base_path = os.path.splitext(self.video_path.get())[0]
output_path = f"{base_path}{suffix}.srt"
# Salva arquivo
with open(output_path, 'w', encoding='utf-8') as f:
f.write(self.subtitle_text.get('1.0', tk.END))
messagebox.showinfo("Sucesso", f"Legendas salvas em:\n{output_path}")
except Exception as e:
messagebox.showerror("Erro", f"Erro ao salvar legendas: {str(e)}")
if __name__ == "__main__":
if check_dependencies():
root = tk.Tk()
app = VideoSubtitleApp(root)
root.mainloop()