File size: 3,398 Bytes
ae55e39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import os
import torch
import librosa
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2FeatureExtractor
# Define emotion labels for the model
EMOTION_LABELS = [
"Neutral", "Happy", "Sad", "Angry", "Fearful", "Disgusted", "Surprised"
]
# Model paths
MODEL_NAME = "superb/wav2vec2-base-superb-er"
# Look for model files directly in the current directory
LOCAL_MODEL_DIR = "."
LOCAL_FEATURE_EXTRACTOR_DIR = "."
def load_model():
"""Load the emotion recognition model and feature extractor"""
try:
# Check if model files exist in current directory
model_files_exist = any(f.startswith("pytorch_model") for f in os.listdir(LOCAL_MODEL_DIR))
config_file_exists = os.path.exists(os.path.join(LOCAL_MODEL_DIR, "config.json"))
feature_extractor_exists = os.path.exists(os.path.join(LOCAL_FEATURE_EXTRACTOR_DIR, "preprocessor_config.json"))
if model_files_exist and config_file_exists and feature_extractor_exists:
print("Loading model and feature extractor from current directory...")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(LOCAL_FEATURE_EXTRACTOR_DIR)
model = Wav2Vec2ForSequenceClassification.from_pretrained(LOCAL_MODEL_DIR)
else:
print("Local model files not found. Loading from Hugging Face...")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_NAME)
model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_NAME)
return model, feature_extractor
except Exception as e:
print(f"Error loading model: {e}")
# Fallback to using Auto classes if specific classes fail
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
return model, feature_extractor
def predict_emotion(audio_path):
"""Predict emotion from audio file"""
try:
# Load model and feature extractor
model, feature_extractor = load_model()
# Load and preprocess audio
speech_array, sampling_rate = librosa.load(audio_path, sr=16000)
# Process the audio
inputs = feature_extractor(speech_array, sampling_rate=sampling_rate, return_tensors="pt")
# Predict emotion
with torch.no_grad():
logits = model(**inputs).logits
# Get emotion label
predicted_class_id = torch.argmax(logits, dim=-1).item()
# Return the predicted emotion
return EMOTION_LABELS[predicted_class_id]
except Exception as e:
print(f"Error predicting emotion: {e}")
return "Error: Could not predict emotion"
# For testing
if __name__ == "__main__":
# Test with a file from current directory
current_dir = "."
wav_files = [f for f in os.listdir(current_dir) if f.endswith(".wav")]
if wav_files:
test_file = wav_files[0]
print(f"Testing with file: {test_file}")
emotion = predict_emotion(test_file)
print(f"Predicted emotion: {emotion}")
else:
print("No .wav files found in current directory for testing.")
|