immelstorun's picture
Update app.py
08a22c7
from speechbrain.pretrained.interfaces import foreign_class
import gradio as gr
import os
import warnings
warnings.filterwarnings("ignore")
# Function to get the list of audio files in the 'rec/' directory
def get_audio_files_list(directory="rec"):
try:
return [f for f in os.listdir(directory) if os.path.isfile(os.path.join(directory, f))]
except FileNotFoundError:
print("The 'rec' directory does not exist. Please make sure it is the correct path.")
return []
# 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(selected_audio):
file_path = os.path.join("rec", selected_audio)
out_prob, score, index, text_lab = learner.classify_file(file_path)
emotion = emotion_dict[text_lab[0]]
return emotion, file_path # Return both emotion and file path
# Get the list of audio files for the dropdown
audio_files_list = get_audio_files_list()
# Loading Gradio interface
inputs = gr.Dropdown(label="Select Audio", choices=audio_files_list)
outputs = [gr.outputs.Textbox(label="Predicted Emotion"), gr.outputs.Audio(label="Play Audio")]
title = "ML Speech Emotion Detection"
description = "Speechbrain powered wav2vec 2.0 pretrained model on IEMOCAP dataset using Gradio."
interface = gr.Interface(fn=predict_emotion, inputs=inputs, outputs=outputs, title=title, description=description)
interface.launch()