godax84/criar_legendas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Modificar main.py

Browse Source
This commit is contained in:
godax84
2024-11-26 10:20:07 -08:00
parent d23aadcd0e
commit ef1798dcf3
1 changed files with 173 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
main.py
View File

@ -7,21 +7,49 @@ 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 de Legendas")
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
# Dicionário de línguas para extração
self.languages = {
'Português (Brasil)': 'pt',
'Português (Portugal)': 'pt',
@ -32,25 +60,50 @@ class VideoSubtitleApp:
'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 modelo em thread separada
# Inicializar modelos em threads separadas
self.model_ready = False
threading.Thread(target=self.initialize_whisper, daemon=True).start()
self.translator = None
threading.Thread(target=self.initialize_models, daemon=True).start()
def initialize_whisper(self):
"""Inicializa o modelo Whisper"""
def initialize_models(self):
"""Inicializa os modelos de Whisper e Tradução"""
try:
self.status_var.set("Carregando modelo...")
# Usar modelo base para melhor equilíbrio
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("Modelo carregado com sucesso")
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 modelo")
messagebox.showerror("Erro", f"Erro ao carregar modelo Whisper:\n{str(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"""
@ -71,7 +124,7 @@ class VideoSubtitleApp:
ttk.Button(top_frame, text="📂 Selecionar Vídeo",
command=self.select_file).pack(side=tk.LEFT, padx=5)
ttk.Label(top_frame, text="🌐 Idioma:").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()),
@ -80,6 +133,15 @@ class VideoSubtitleApp:
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))
@ -100,6 +162,11 @@ class VideoSubtitleApp:
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')
@ -130,6 +197,81 @@ class VideoSubtitleApp:
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(
@ -246,7 +388,7 @@ class VideoSubtitleApp:
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"
@ -274,6 +416,7 @@ class VideoSubtitleApp:
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)
@ -285,18 +428,31 @@ class VideoSubtitleApp:
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:
output_path = os.path.splitext(self.video_path.get())[0] + ".srt"
# 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__":
root = tk.Tk()
app = VideoSubtitleApp(root)
root.mainloop()
if check_dependencies():
root = tk.Tk()
app = VideoSubtitleApp(root)
root.mainloop()