Translation / app.py
JoelJokku's picture
Updated codes
e0f07e3
import whisper
import numpy as np
from datetime import datetime
import torch
from flask import Flask,request,jsonify
from transformers import pipeline
import os
from dotenv import load_dotenv
from flask_cors import CORS
# Add this before loading the model
os.environ["XDG_CACHE_HOME"] = os.getcwd() + "/.cache" # Set custom cache location
load_dotenv()
# Then load the model
app=Flask(__name__)
CORS(app)
cache_dir = os.path.join(os.getcwd(), ".cache")
os.makedirs(cache_dir, exist_ok=True)
model = whisper.load_model("base", download_root=cache_dir) # Explicit cache location
SAMPLE_RATE=16000
DURATION=10
OUTPUT_FILE="recorded_audio.wav"
device="cuda" if torch.cuda.is_available() else "cpu"
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-mul-en",token=os.environ.get("HUGGING_FACE_AUTH_TOKEN"))
def detect(audio_path):
audio=whisper.load_audio(audio_path)
audio=whisper.pad_or_trim(audio)
mel=whisper.log_mel_spectrogram(audio).to(model.device)
_,probs=model.detect_language(mel)
detected_lang=max(probs,key=probs.get)
conf=probs[detected_lang]
res=model.transcribe(audio_path,language=detected_lang)
trans=res["text"]
translation = trans if detected_lang == "en" else translator(trans)[0]["translation_text"]
return detected_lang,conf,trans,translation
@app.route('/detect-language',methods=['POST'])
def detect_endpoint():
try:
if 'audio' not in request.files:
return jsonify({"status":"error","message":"No audio file provided"}),400
else:
audio_file=request.files['audio']
audio_path=os.path.join("uploads",audio_file.filename)
os.makedirs("uploads",exist_ok=True)
audio_file.save(audio_path)
'''else:
audio_path=record()'''
det_lang,conf,_,_=detect(audio_path)
if os.path.exists(audio_path) and audio_path!=OUTPUT_FILE:
os.remove(audio_path)
return jsonify({
"status":"success",
"detected_language":det_lang,
"confidence":f"{conf:.2%}"
})
except Exception as e:
return jsonify({"status":"error","message":str(e)}),500
@app.route('/transcription',methods=['POST'])
def transcription_endpoint():
try:
if 'audio' not in request.files:
return jsonify({"status":"error","message":"No audio file provided"}),400
else:
audio_file=request.files['audio']
audio_path=os.path.join("uploads",audio_file.filename)
os.makedirs("uploads",exist_ok=True)
audio_file.save(audio_path)
'''else:
audio_path=record()'''
_,_,trans,_=detect(audio_path)
if os.path.exists(audio_path) and audio_path!=OUTPUT_FILE:
os.remove(audio_path)
return jsonify({
"status":"success",
"transcription":trans
})
except Exception as e:
return jsonify({"status":"error","message":str(e)}),500
@app.route('/translation',methods=['POST'])
def translation_endpoint():
try:
if 'audio' not in request.files:
return jsonify({"status":"error","message":"No audio file provided"}),400
else:
audio_file=request.files['audio']
audio_path=os.path.join("uploads",audio_file.filename)
os.makedirs("uploads",exist_ok=True)
audio_file.save(audio_path)
'''else:
audio_path=record()'''
det_lang,_,_,trans=detect(audio_path)
if os.path.exists(audio_path) and audio_path!=OUTPUT_FILE:
os.remove(audio_path)
return jsonify({
"status":"success",
"detected_language":det_lang,
"translation":trans
})
except Exception as e:
return jsonify({"status":"error","message":str(e)}),500
if __name__=="__main__":
app.run(debug=False,host="0.0.0.0",port=int(os.environ.get("PORT",7860)))