Spaces:
Runtime error
Runtime error
File size: 2,035 Bytes
ebf88d6 3a70449 ebf88d6 3a70449 ebf88d6 3a70449 |
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 |
import librosa
import torch
import numpy as np
def preprocess_audio(audio_array, feature_extractor, sampling_rate, max_length=3000):
"""
Preprocesses audio for emotion prediction.
"""
inputs = feature_extractor(
audio_array,
sampling_rate=sampling_rate,
return_tensors="pt",
)
mel_features = inputs["input_features"]
current_length = mel_features.size(2)
if current_length < max_length:
pad_size = max_length - current_length
mel_features = torch.nn.functional.pad(mel_features, (0, pad_size), mode="constant", value=0)
elif current_length > max_length:
mel_features = mel_features[:, :, :max_length]
inputs["input_features"] = mel_features
return inputs
def predict_emotion(audio_path, model, feature_extractor, id2label, sampling_rate=16000, chunk_duration=8.0):
"""
Predicts emotions from an audio file.
"""
audio_array, _ = librosa.load(audio_path, sr=sampling_rate)
chunk_length = int(sampling_rate * chunk_duration)
num_chunks = len(audio_array) // chunk_length + int(len(audio_array) % chunk_length > 0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
results = []
for i in range(num_chunks):
start = i * chunk_length
end = min((i + 1) * chunk_length, len(audio_array))
chunk = audio_array[start:end]
start_time = round(start / sampling_rate, 2)
end_time = round(end / sampling_rate, 2)
inputs = preprocess_audio(chunk, feature_extractor, sampling_rate, max_length=3000)
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
predicted_id = torch.argmax(logits, dim=-1).item()
predicted_label = id2label[predicted_id]
results.append({"chunk": i + 1, "start_time": start_time, "end_time": end_time, "emotion": predicted_label})
return results
|