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(file_path):
    if file_path:
        # Classify the file
        out_prob, score, index, text_lab = classifier.classify_file(file_path)
        if isinstance(text_lab, list):
            text_lab = text_lab[0]
        # Map the original labels to the desired categories
        emotion_mapping = {
            'neu': 'Neutral',
            'ang': 'Angry',
            'hap': 'Happy',
            'sad': 'Sadness'
        }
        # Get the corresponding category from the mapping
        emotion_category = emotion_mapping.get(text_lab, 'Unknown')
        
        emotion_category = emotion_mapping.get(text_lab, 'Unknown')
        # Return the emotion category
        st.write(emotion_category)
    else:
        st.write("Please provide the path to an audio file.")


def main():
    st.title("Emotion Recognition")
    file_path = st.text_input("Enter the path of the audio file (e.g., /path/to/audio.wav):")
    if file_path:
        emotion(file_path)

if __name__ == "__main__":
    main()