import streamlit as st
import tempfile
import os
from speechbrain.inference.interfaces import foreign_class

# Initialize the classifier
classifier = foreign_class(source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP", pymodule_file="custom_interface.py", classname="CustomEncoderWav2vec2Classifier")

def save_uploaded_file(uploaded_file):
    temp_dir = tempfile.TemporaryDirectory()
    file_path = os.path.join(temp_dir.name, uploaded_file.name)
    with open(file_path, "wb") as f:
        f.write(uploaded_file.getbuffer())
    return file_path

def emotion(uploaded_file):
    if uploaded_file is not None:
        # Save the uploaded file to a temporary location
        file_path = save_uploaded_file(uploaded_file)
        # Classify the file
        out_prob, score, index, text_lab = classifier.classify_file(file_path)
        # Display the output
        st.write(text_lab)
    else:
        st.write("Please upload a file.")

def main():
    st.title("Emotion Recognition")
    uploaded_file = st.file_uploader("Upload audio file", type=["wav"])
    if uploaded_file is not None:
        emotion(uploaded_file)

if __name__ == "__main__":
    main()