Update src/main.py

src/main.py

@@ -1,47 +1,57 @@
 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
-from transformers import MarianMTModel, MarianTokenizer
-import torch
+from flask import Flask, render_template, Response, request, send_file
+import io
+import cv2
+import numpy as np
 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 = MarianTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-model = MarianMTModel.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()
 
 def translate_korean_to_english(text):
    try:
-       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}")
-           return translation
-       return text
+       url = "https://translate.googleapis.com/translate_a/single"
+       params = {
+           "client": "gtx",
+           "sl": "ko",
+           "tl": "en",
+           "dt": "t",
+           "q": text
+       }
+       response = requests.get(url, params=params)
+       return response.json()[0][0][0]
    except Exception as e:
        print(f"Translation error: {e}")
        return text
 
+def generate_complete_video(gloss_list, dataset, list_2000_tokens):
+   frames = []
+   for frame in dg.generate_video(gloss_list, dataset, list_2000_tokens):
+       frame_data = frame.split(b'\r\n\r\n')[1]
+       nparr = np.frombuffer(frame_data, np.uint8)
+       img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+       frames.append(img)
+   
+   height, width = frames[0].shape[:2]
+   fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+   temp_path = os.path.join('/tmp', 'temp.mp4')
+   out = cv2.VideoWriter(temp_path, fourcc, 25, (width, height))
+   
+   for frame in frames:
+       out.write(frame)
+   out.release()
+   
+   with open(temp_path, 'rb') as f:
+       video_bytes = f.read()
+   
+   os.remove(temp_path)
+   return video_bytes
+
 @app.route('/')
 def index():
    return render_template('index.html', title=app.config['TITLE'])
@@ -52,10 +62,6 @@ def result():
        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()
            
@@ -73,7 +79,6 @@ def result():
                                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')
@@ -83,5 +88,16 @@ def video_feed():
    return Response(dg.generate_video(gloss_list, dataset, list_2000_tokens), 
                   mimetype='multipart/x-mixed-replace; boundary=frame')
 
+@app.route('/download_video/<gloss_sentence>')
+def download_video(gloss_sentence):
+   gloss_list = gloss_sentence.split()
+   video_bytes = generate_complete_video(gloss_list, dataset, list_2000_tokens)
+   return send_file(
+       io.BytesIO(video_bytes),
+       mimetype='video/mp4',
+       as_attachment=True,
+       download_name='sign_language.mp4'
+   )
+
 if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)