TransLingo / app.py
mishiawan's picture
Update app.py
0e542a8 verified
import streamlit as st
import whisper
from transformers import pipeline
import speech_recognition as sr
from gtts import gTTS
import tempfile
import os
import base64
# Load Whisper model for Speech-to-Text
@st.cache_resource
def load_whisper_model():
return whisper.load_model("base") # You can use other models like "large" for better accuracy
# Load translation model from Hugging Face
@st.cache_resource
def load_translation_model(input_lang, output_lang):
model_name = f"Helsinki-NLP/opus-mt-{input_lang}-{output_lang}"
return pipeline("translation", model=model_name)
# Initialize speech recognition
recognizer = sr.Recognizer()
def speech_to_text(audio_file):
with sr.AudioFile(audio_file) as source:
audio_data = recognizer.record(source)
return recognizer.recognize_google(audio_data)
def translate_text(input_text, translation_pipeline):
return translation_pipeline(input_text)[0]["translation_text"]
def text_to_speech(text, language):
tts = gTTS(text=text, lang=language)
temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
tts.save(temp_file.name)
return temp_file.name
# Streamlit UI
st.title("Voice-to-Voice Translator")
# Sidebar for user input
st.sidebar.header("Settings")
input_lang = st.sidebar.selectbox("Select Input Language", ["en", "fr", "es", "de", "hi"])
output_lang = st.sidebar.selectbox("Select Output Language", ["en", "fr", "es", "de", "hi"])
# Initialize models
whisper_model = load_whisper_model()
translation_pipeline = load_translation_model(input_lang, output_lang)
# Input options
st.header("Input your voice or text")
# Option to upload an audio file
audio_file = st.file_uploader("Upload an audio file (WAV format)", type=["wav"])
if audio_file is not None:
# Convert speech to text
input_text = speech_to_text(audio_file)
st.write(f"Recognized Text: {input_text}")
# Translate the text
translated_text = translate_text(input_text, translation_pipeline)
st.success(f"Translated Text: {translated_text}")
# Convert translated text to speech
output_audio = text_to_speech(translated_text, output_lang)
audio_bytes = open(output_audio, "rb").read()
# Play the audio
st.audio(audio_bytes, format="audio/mp3")
# Provide download link for audio
b64 = base64.b64encode(audio_bytes).decode()
href = f'<a href="data:audio/mp3;base64,{b64}" download="translated_audio.mp3">Download Translated Audio</a>'
st.markdown(href, unsafe_allow_html=True)
else:
st.write("Please upload an audio file to get started.")