evolve-build/visual_emotion_recognition

1import time
2import numpy as np
3import torch
4import torch.nn.functional as F
5import librosa
6import pyaudio
7from collections import deque
8import threading
9
10# Device configuration
11device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
12
13####################################
14# Load Pretrained TorchScript Model
15####################################
16MODEL_TS_PATH = "complex_audio_emotion_model_ts.pt"  # Ensure this file is in your working directory
17model = torch.jit.load(MODEL_TS_PATH, map_location=device)
18model.eval()
19
20####################################
21# Inverse Label Mapping (must match your training)
22####################################
23# Example mapping if you trained with 8 classes:
24idx_to_emotion = {
25    0: "angry",
26    1: "calm",
27    2: "happy",
28    3: "sad",
29    4: "fearful",
30    5: "disgust",
31    6: "surprised",
32    7: "neutral"
33}
34
35####################################
36# Hardcoded Nervousness Mapping Function
37####################################
38base_nervousness = {
39    "angry": 80,
40    "fearful": 90,
41    "disgust": 70,
42    "sad": 60,
43    "surprised": 55,
44    "happy": 20,
45    "neutral": 40,
46    "calm": 30
47}
48
49def compute_nervousness(predicted_emotion, f0_std):
50    """
51    Computes a nervousness score (0-100) using:
52      - A base score defined by the predicted emotion.
53      - An adjustment based on the standard deviation of the fundamental frequency (f0_std):
54            f0_std < 20 Hz: +0,
55            20 Hz <= f0_std < 50 Hz: +10,
56            f0_std >= 50 Hz: +20.
57    """
58    base_score = base_nervousness.get(predicted_emotion, 40)
59    if f0_std < 20:
60        adjustment = 0
61    elif f0_std < 50:
62        adjustment = 10
63    else:
64        adjustment = 20
65    final_score = base_score + adjustment
66    return np.clip(final_score, 0, 100)
67
68####################################
69# PyAudio Setup for Real-Time Audio Capture
70####################################
71AUDIO_RATE = 16000  # Sampling rate (Hz)
72CHUNK = 1024       # Number of samples per frame
73audio_buffer = deque(maxlen=AUDIO_RATE)  # 1-second buffer
74buffer_lock = threading.Lock()
75
76def audio_callback(in_data, frame_count, time_info, status):
77    # Convert the incoming byte data to a numpy array of float32 normalized to [-1, 1]
78    data = np.frombuffer(in_data, dtype=np.int16).astype(np.float32) / 32768.0
79    with buffer_lock:
80        audio_buffer.extend(data.tolist())
81    return (in_data, pyaudio.paContinue)
82
83def start_audio_stream():
84    p = pyaudio.PyAudio()
85    stream = p.open(format=pyaudio.paInt16,
86                    channels=1,
87                    rate=AUDIO_RATE,
88                    input=True,
89                    frames_per_buffer=CHUNK,
90                    stream_callback=audio_callback)
91    stream.start_stream()
92    return p, stream
93
94####################################
95# Real-Time Inference Function
96####################################
97def real_time_inference(update_threshold=5.0, smoothing_window=5):
98    """
99    Captures 1-second audio windows from the microphone and computes:
100      - The predicted emotion from the transformer-based classifier.
101      - The nervousness score based on the predicted emotion and pitch variability.
102    The output is smoothed over a sliding window and only updated when a significant change occurs.
103    """
104    p, stream = start_audio_stream()
105    print("Real-time inference started. Press Ctrl+C to stop.")
106    
107    recent_scores = deque(maxlen=smoothing_window)
108    last_emotion = None
109    last_nervousness = None
110
111    try:
112        while True:
113            time.sleep(1)  # Process every second
114            with buffer_lock:
115                if len(audio_buffer) < AUDIO_RATE:
116                    continue  # Wait until we have 1 second of audio
117                current_audio = np.array(list(audio_buffer))[:AUDIO_RATE]
118            # Prepare input tensor [1, num_samples]
119            audio_tensor = torch.tensor(current_audio, dtype=torch.float32).unsqueeze(0).to(device)
120            
121            # Run inference on audio through the model
122            with torch.no_grad():
123                logits, _ = model(audio_tensor)
124                probs = F.softmax(logits, dim=1)
125                pred_idx = torch.argmax(probs, dim=1).item()
126            predicted_emotion = idx_to_emotion.get(pred_idx, "neutral")
127            
128            # Compute fundamental frequency using librosa.pyin
129            f0, _, _ = librosa.pyin(current_audio, fmin=librosa.note_to_hz('C2'),
130                                    fmax=librosa.note_to_hz('C7'))
131            f0 = f0[~np.isnan(f0)]
132            f0_std = np.std(f0) if len(f0) > 0 else 0.0
133            
134            # Compute nervousness score
135            current_score = compute_nervousness(predicted_emotion, f0_std)
136            recent_scores.append(current_score)
137            smoothed_score = np.mean(recent_scores)
138            
139            # Update output if emotion changes or if nervousness change exceeds threshold
140            if (last_emotion != predicted_emotion) or (last_nervousness is None or abs(smoothed_score - last_nervousness) >= update_threshold):
141                print(f"Emotion: {predicted_emotion} | Nervousness: {smoothed_score:.1f}/100")
142                last_emotion = predicted_emotion
143                last_nervousness = smoothed_score
144
145    except KeyboardInterrupt:
146        print("Real-time inference stopped.")
147    finally:
148        stream.stop_stream()
149        stream.close()
150        p.terminate()
151
152if __name__ == "__main__":
153    real_time_inference()
154