emotion / src /emotion_detector.py
jc047's picture
Update src/emotion_detector.py
ea8ce67 verified
raw
history blame
811 Bytes
# emotion_detector.py
import torchaudio
from speechbrain.inference.classifiers import AudioClassifier
import os
import tempfile
# Use a temp directory that's guaranteed to be writable
temp_dir = os.path.join(tempfile.gettempdir(), "emotion_model")
# Load pretrained model into temp directory
classifier = AudioClassifier.from_hparams(
source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
savedir=temp_dir
)
EMOTION_EMOJIS = {
"angry": "😠",
"happy": "πŸ˜„",
"neutral": "😐",
"sad": "😒",
"fearful": "😨"
}
def detect_emotion(audio_path):
signal, fs = torchaudio.load(audio_path)
prediction = classifier.classify_file(audio_path)
emotion = prediction[3] # Predicted label
emoji = EMOTION_EMOJIS.get(emotion.lower(), "❓")
return emotion, emoji