From ef1798dcf392f54f05815cc9422cb5f21c1e94f4 Mon Sep 17 00:00:00 2001 From: godax84 Date: Tue, 26 Nov 2024 10:20:07 -0800 Subject: [PATCH] Modificar main.py --- main.py | 190 +++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 173 insertions(+), 17 deletions(-) diff --git a/main.py b/main.py index 7ce61c5..2d9e13b 100644 --- a/main.py +++ b/main.py @@ -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() \ No newline at end of file + if check_dependencies(): + root = tk.Tk() + app = VideoSubtitleApp(root) + root.mainloop() + \ No newline at end of file