Modificar main.py

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):
+    def initialize_models(self):
-        """Inicializa o modelo Whisper"""
+        """Inicializa os modelos de Whisper e Tradução"""
         try:
-            self.status_var.set("Carregando modelo...")
+            self.status_var.set("Carregando modelos...")
-            # Usar modelo base para melhor equilíbrio
+            # 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")
+            self.status_var.set("Erro ao carregar modelos")
-            messagebox.showerror("Erro", f"Erro ao carregar modelo Whisper:\n{str(e)}")
+            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()
+    if check_dependencies():
-    app = VideoSubtitleApp(root)
+        root = tk.Tk()
-    root.mainloop()
+        app = VideoSubtitleApp(root)
+        root.mainloop()
+