Spaces:
Runtime error
Runtime error
| import torch | |
| import librosa | |
| import numpy as np | |
| from transformers import AutoFeatureExtractor, Wav2Vec2ForSequenceClassification | |
| import argparse | |
| import os | |
| class EmotionPredictor: | |
| def __init__(self, model_path="models/wav2vec2-finetuned"): | |
| """ | |
| Initializes the predictor with the fine-tuned model. | |
| """ | |
| print(f"Loading model from: {model_path}...") | |
| # Load config and model | |
| try: | |
| self.feature_extractor = AutoFeatureExtractor.from_pretrained(model_path) | |
| self.model = Wav2Vec2ForSequenceClassification.from_pretrained(model_path) | |
| self.model.eval() # Set to inference mode | |
| # Extract labels from config | |
| self.id2label = self.model.config.id2label | |
| print(f"Model loaded. Labels: {list(self.id2label.values())}") | |
| except Exception as e: | |
| print(f"Failed to load model: {e}") | |
| raise | |
| def predict(self, audio_path): | |
| """ | |
| Predicts emotion from a single audio file. | |
| """ | |
| # 1. Load Audio | |
| try: | |
| speech, sr = librosa.load(audio_path, sr=16000) | |
| except Exception as e: | |
| return {"error": f"Could not load audio: {e}"} | |
| # 2. Preprocess | |
| inputs = self.feature_extractor( | |
| speech, | |
| sampling_rate=16000, | |
| padding=True, | |
| return_tensors="pt" | |
| ) | |
| # 3. Inference | |
| with torch.no_grad(): | |
| logits = self.model(inputs.input_values).logits | |
| # 4. Post-processing | |
| probabilities = torch.nn.functional.softmax(logits, dim=-1) | |
| predicted_id = torch.argmax(logits, dim=-1).item() | |
| confidence = probabilities[0][predicted_id].item() | |
| predicted_label = self.id2label[predicted_id] | |
| return { | |
| "emotion": predicted_label, | |
| "confidence": f"{confidence:.2%}", | |
| "probabilities": { | |
| self.id2label[i]: f"{probabilities[0][i].item():.2%}" | |
| for i in range(len(self.id2label)) | |
| } | |
| } | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Predict emotion from audio file") | |
| parser.add_argument("--file", type=str, help="Path to .wav file") | |
| args = parser.parse_args() | |
| if args.file: | |
| predictor = EmotionPredictor() | |
| result = predictor.predict(args.file) | |
| print("\nPrediction Result:") | |
| print(f"Emotion: {result.get('emotion', 'Error')}") | |
| print(f"Confidence: {result.get('confidence', 'N/A')}") | |
| print(result) | |
| else: | |
| print("Please provide a file path using --file") | |