from speechbrain.pretrained.interfaces import foreign_class

import warnings
warnings.filterwarnings("ignore")

import os
import gradio as gr

# Путь к каталогу с предзаписанными аудиофайлами
prerecorded_audio_path = 'prerecorded'
# Список файлов в каталоге prerecorded
prerecorded_audio_files = os.listdir(prerecorded_audio_path)
# Полные пути к файлам для Dropdown
prerecorded_audio_files_full_path = [os.path.join(prerecorded_audio_path, file) for file in prerecorded_audio_files]

# Loading the speechbrain emotion detection model
learner = foreign_class(
    source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
    pymodule_file="custom_interface.py",
    classname="CustomEncoderWav2vec2Classifier"
)

# Building prediction function for gradio
emotion_dict = {
    'sad': 'Sad',
    'hap': 'Happy',
    'ang': 'Anger',
    'fea': 'Fear',
    'sur': 'Surprised',
    'neu': 'Neutral'
}

def predict_emotion(uploaded_audio=None, prerecorded_audio=None):
    # Если выбран аудиофайл из выпадающего списка, использовать его
    if prerecorded_audio is not None:
        audio_file_path = prerecorded_audio
    elif uploaded_audio is not None:
        # Иначе, если загружен файл, использовать его
        audio_file_path = uploaded_audio.name
    else:
        # Если нет файла, вернуть сообщение об ошибке
        return "No audio file provided", 0

    out_prob, score, index, text_lab = learner.classify_file(audio_file_path)
    emotion_probability = out_prob[0][index[0]].item()
    
    # Возвращаем словарь с эмоцией и вероятностью
    return {"Emotion": emotion_dict[text_lab[0]], "Probability": f"{emotion_probability:.2f}"}

# Модифицированный Gradio interface
inputs = [
    gr.inputs.Dropdown(list(prerecorded_audio_files_full_path), label="Select Prerecorded Audio", default=None),
    gr.inputs.Audio(label="Or Upload Audio", type="file", source="upload", optional=True),
    gr.inputs.Audio(label="Or Record Audio", type="file", source="microphone", optional=True)
]

outputs = gr.outputs.Label(num_top_classes=2)

title = "ML Speech Emotion Detection"
description = "Detect emotions from speech using a Speechbrain powered model."

gr.Interface(fn=predict_emotion, inputs=inputs, outputs=outputs, title=title, description=description).launch()