Spaces:
Sleeping
Sleeping
File size: 916 Bytes
3461076 | 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 | import torch
from model import CLSTMModel
from config import CONFIG, EMOTION_CONFIG
from audio_utils import preprocess_audio
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = torch.load(CONFIG["model_path"], map_location=device)
if "label_map" in checkpoint:
inv = {v: k for k, v in checkpoint["label_map"].items()}
emotions = [inv[i] for i in range(len(inv))]
else:
emotions = list(EMOTION_CONFIG.keys())
model = CLSTMModel(
n_mels=CONFIG["n_mels"],
n_classes=len(emotions)
).to(device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
def predict(path):
x = preprocess_audio(path, device)
with torch.no_grad():
logits = model(x)
probs = torch.softmax(logits, dim=1)
idx = torch.argmax(probs, dim=1).item()
return {
"emotion": emotions[idx],
"confidence": float(probs[0][idx])
}
|