Update src/main.py

src/main.py

@@ -2,81 +2,91 @@ import display_gloss as dg
 import synonyms_preprocess as sp
 from NLP_Spacy_base_translator import NlpSpacyBaseTranslator 
 from flask import Flask, render_template, Response, request
-import googletrans
-from googletrans import Translator
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqTranslation
+import torch
+import os
 
 app = Flask(__name__, static_folder='static')
 app.config['TITLE'] = 'Sign Language Translate'
 
+# Set cache directory
+cache_dir = "/tmp/huggingface"
+if not os.path.exists(cache_dir):
+   os.makedirs(cache_dir, exist_ok=True)
+os.environ['TRANSFORMERS_CACHE'] = cache_dir
+os.environ['HF_HOME'] = cache_dir
+
+# Force CPU usage
+device = torch.device('cpu')
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+
+# Load pre-trained Korean-English translation model
+model_name = "Helsinki-NLP/opus-mt-ko-en"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+model = AutoModelForSeq2SeqTranslation.from_pretrained(model_name, cache_dir=cache_dir)
+model = model.to(device)
+
 nlp, dict_docs_spacy = sp.load_spacy_values()
 dataset, list_2000_tokens = dg.load_data()
-translator = Translator(service_urls=['translate.google.com'])
 
 def translate_korean_to_english(text):
-    try:
-        # 한글 감지
-        if any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in text):
-            translation = translator.translate(text, src='ko', dest='en')
-            print(f"Translation result: {translation.text}")  # 디버깅용
-            return translation.text
-        return text
-    except Exception as e:
-        print(f"Translation error: {e}")
-        try:
-            # 백업 방식으로 파파고 API 사용
-            headers = {
-                'X-Naver-Client-Id': 'YOUR_CLIENT_ID',
-                'X-Naver-Client-Secret': 'YOUR_CLIENT_SECRET'
-            }
-            data = {
-                'source': 'ko',
-                'target': 'en',
-                'text': text
-            }
-            response = requests.post('https://openapi.naver.com/v1/papago/n2mt', headers=headers, data=data)
-            return response.json()['message']['result']['translatedText']
-        except:
-            return text
+   try:
+       # Check if input is Korean
+       if any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in text):
+           inputs = tokenizer(text, return_tensors="pt", padding=True)
+           outputs = model.generate(**inputs)
+           translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
+           print(f"Translated text: {translation}")  # Debug log
+           return translation
+       return text
+   except Exception as e:
+       print(f"Translation error: {e}")
+       return text
 
 @app.route('/')
 def index():
-    return render_template('index.html', title=app.config['TITLE'])
+   return render_template('index.html', title=app.config['TITLE'])
 
 @app.route('/translate/', methods=['POST'])
 def result():
-    if request.method == 'POST':
-        input_text = request.form['inputSentence']
-        try:
-            english_text = translate_korean_to_english(input_text)
-            if english_text == input_text and any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in input_text):
-                raise Exception("Translation failed")
-                
-            eng_to_asl_translator = NlpSpacyBaseTranslator(sentence=english_text)
-            generated_gloss = eng_to_asl_translator.translate_to_gloss()
-            
-            gloss_list_lower = [gloss.lower() for gloss in generated_gloss.split() if gloss.isalnum()]
-            gloss_sentence_before_synonym = " ".join(gloss_list_lower)
-            
-            gloss_list = [sp.find_synonyms(gloss, nlp, dict_docs_spacy, list_2000_tokens) 
-                         for gloss in gloss_list_lower]
-            gloss_sentence_after_synonym = " ".join(gloss_list)
-            
-            return render_template('result.html',
-                                title=app.config['TITLE'],
-                                original_sentence=input_text,
-                                english_translation=english_text,
-                                gloss_sentence_before_synonym=gloss_sentence_before_synonym,
-                                gloss_sentence_after_synonym=gloss_sentence_after_synonym)
-        except Exception as e:
-            print(f"Error in translation process: {e}")
-            return render_template('error.html', error=str(e))
+   if request.method == 'POST':
+       input_text = request.form['inputSentence']
+       try:
+           # Translate to English
+           english_text = translate_korean_to_english(input_text)
+           
+           # Check if translation failed
+           if english_text == input_text and any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in input_text):
+               raise Exception("Translation failed")
+           
+           # Convert to ASL gloss
+           eng_to_asl_translator = NlpSpacyBaseTranslator(sentence=english_text)
+           generated_gloss = eng_to_asl_translator.translate_to_gloss()
+           
+           # Process gloss
+           gloss_list_lower = [gloss.lower() for gloss in generated_gloss.split() if gloss.isalnum()]
+           gloss_sentence_before_synonym = " ".join(gloss_list_lower)
+           
+           gloss_list = [sp.find_synonyms(gloss, nlp, dict_docs_spacy, list_2000_tokens) 
+                        for gloss in gloss_list_lower]
+           gloss_sentence_after_synonym = " ".join(gloss_list)
+           
+           return render_template('result.html',
+                               title=app.config['TITLE'],
+                               original_sentence=input_text,
+                               english_translation=english_text,
+                               gloss_sentence_before_synonym=gloss_sentence_before_synonym,
+                               gloss_sentence_after_synonym=gloss_sentence_after_synonym)
+       except Exception as e:
+           print(f"Error in translation process: {e}")
+           return render_template('error.html', error=str(e))
 
 @app.route('/video_feed')
 def video_feed():
-    sentence = request.args.get('gloss_sentence_to_display', '')
-    gloss_list = sentence.split()
-    return Response(dg.generate_video(gloss_list, dataset, list_2000_tokens), 
-                   mimetype='multipart/x-mixed-replace; boundary=frame')
+   sentence = request.args.get('gloss_sentence_to_display', '')
+   gloss_list = sentence.split()
+   return Response(dg.generate_video(gloss_list, dataset, list_2000_tokens), 
+                  mimetype='multipart/x-mixed-replace; boundary=frame')
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=True)
+   app.run(host="0.0.0.0", port=7860, debug=True)