evolve-build/visual_emotion_recognition

1import cv2
2import time
3import numpy as np
4import torch
5import torch.nn.functional as F
6import mediapipe as mp
7from torchvision import transforms
8from collections import deque
9from PIL import Image  # Added missing import
10
11# Device configuration
12device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
14####################################
15# Load Pretrained TorchScript Visual Model
16####################################
17MODEL_TS_PATH = "complex_visual_emotion_model_ts.pt"  # Ensure this file is in your working directory
18model = torch.jit.load(MODEL_TS_PATH, map_location=device)
19model.eval()
20
21####################################
22# Inverse Label Mapping (should match training)
23####################################
24# For example, if you trained with FER2013 (7 classes):
25idx_to_emotion = {
26    0: "angry",
27    1: "disgust",
28    2: "fear",
29    3: "happy",
30    4: "sad",
31    5: "surprise",
32    6: "neutral"
33}
34
35####################################
36# Hardcoded Nervousness Mapping for Visual Cues
37####################################
38# Base nervousness scores for each emotion (values based on literature)
39base_nervousness = {
40    "angry": 80,
41    "fear": 90,
42    "disgust": 70,
43    "sad": 60,
44    "surprise": 55,
45    "happy": 20,
46    "neutral": 40
47}
48
49def compute_nervousness_visual(predicted_emotion, furrow_intensity):
50    """
51    Computes a nervousness score (0-100) using:
52      - A base score for the predicted emotion.
53      - An adjustment based on eyebrow furrow intensity (normalized 0-1).
54    For example:
55        furrow_intensity < 0.3: adjustment = 0,
56        0.3 <= furrow_intensity < 0.6: adjustment = +10,
57        furrow_intensity >= 0.6: adjustment = +20.
58    """
59    base_score = base_nervousness.get(predicted_emotion, 40)
60    if furrow_intensity < 0.3:
61        adjustment = 0
62    elif furrow_intensity < 0.6:
63        adjustment = 10
64    else:
65        adjustment = 20
66    final_score = base_score + adjustment
67    return np.clip(final_score, 0, 100)
68
69####################################
70# Set Up Image Preprocessing for Face Crop
71####################################
72preprocess = transforms.Compose([
73    transforms.Resize((224, 224)),
74    transforms.ToTensor(),
75    transforms.Normalize(mean=[0.485, 0.456, 0.406],
76                         std=[0.229, 0.224, 0.225])
77])
78
79####################################
80# Set Up MediaPipe Face Mesh
81####################################
82mp_face_mesh = mp.solutions.face_mesh
83face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False,
84                                  max_num_faces=1,
85                                  refine_landmarks=True,
86                                  min_detection_confidence=0.5,
87                                  min_tracking_confidence=0.5)
88
89####################################
90# Compute Eyebrow Furrow Intensity
91####################################
92def extract_eyebrow_furrow_intensity(image, landmarks):
93    """
94    Computes a proxy for eyebrow furrow intensity by measuring the Euclidean distance
95    between the inner corners of the eyebrows (landmark indices 70 and 300).
96    The distance is then normalized such that:
97      - Distance >= 50 pixels -> intensity = 0.0 (no furrow)
98      - Distance <= 40 pixels -> intensity = 1.0 (max furrow)
99      - Otherwise, linearly interpolated.
100    """
101    h, w, _ = image.shape
102    try:
103        left_inner = landmarks.landmark[70]
104        right_inner = landmarks.landmark[300]
105    except IndexError:
106        return 0.0
107    x_left, y_left = int(left_inner.x * w), int(left_inner.y * h)
108    x_right, y_right = int(right_inner.x * w), int(right_inner.y * h)
109    distance = np.sqrt((x_right - x_left)**2 + (y_right - y_left)**2)
110    
111    max_distance = 50.0
112    min_distance = 40.0
113    if distance >= max_distance:
114        intensity = 0.0
115    elif distance <= min_distance:
116        intensity = 1.0
117    else:
118        intensity = (max_distance - distance) / (max_distance - min_distance)
119    return intensity
120
121####################################
122# Real-Time Visual Inference Function
123####################################
124def real_time_visual_inference(update_threshold=5.0, smoothing_window=5):
125    cap = cv2.VideoCapture(0)
126    if not cap.isOpened():
127        print("Error: Could not open webcam.")
128        return
129    
130    recent_scores = deque(maxlen=smoothing_window)
131    last_emotion = None
132    last_nervousness = None
133
134    print("Real-time visual inference started. Press 'q' to quit.")
135    
136    while True:
137        ret, frame = cap.read()
138        if not ret:
139            continue
140        
141        # Flip frame for mirror effect and convert BGR to RGB
142        frame = cv2.flip(frame, 1)
143        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
144        
145        # Process the frame with MediaPipe Face Mesh
146        results = face_mesh.process(rgb_frame)
147        
148        display_text = "No face detected"
149        if results.multi_face_landmarks:
150            landmarks = results.multi_face_landmarks[0]
151            
152            # Compute bounding box around the face
153            h, w, _ = frame.shape
154            pts = np.array([(int(lm.x * w), int(lm.y * h)) for lm in landmarks.landmark])
155            x, y, w_box, h_box = cv2.boundingRect(pts)
156            face_crop = frame[y:y+h_box, x:x+w_box]
157            
158            # Preprocess the face crop
159            try:
160                face_img = cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB)
161                face_img = cv2.resize(face_img, (224, 224))
162            except Exception as e:
163                face_img = cv2.resize(frame, (224, 224))
164            face_tensor = preprocess(Image.fromarray(face_img)).unsqueeze(0).to(device)
165            
166            # Run the visual emotion model
167            with torch.no_grad():
168                logits = model(face_tensor)
169                probs = F.softmax(logits, dim=1)
170                pred_idx = torch.argmax(probs, dim=1).item()
171            predicted_emotion = idx_to_emotion.get(pred_idx, "neutral")
172            
173            # Extract eyebrow furrow intensity
174            furrow_intensity = extract_eyebrow_furrow_intensity(frame, landmarks)
175            current_score = compute_nervousness_visual(predicted_emotion, furrow_intensity)
176            recent_scores.append(current_score)
177            smoothed_score = np.mean(recent_scores)
178            
179            display_text = f"Emotion: {predicted_emotion} | Nervousness: {smoothed_score:.1f}/100"
180            last_emotion = predicted_emotion
181            last_nervousness = smoothed_score
182            
183            # Draw face bounding box
184            cv2.rectangle(frame, (x, y), (x+w_box, y+h_box), (0, 255, 0), 2)
185        
186        # Overlay display text on the frame
187        cv2.putText(frame, display_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1,
188                    (0, 255, 0) if "No face" not in display_text else (0, 0, 255), 2)
189        cv2.imshow("Real-Time Visual Emotion & Nervousness", frame)
190        
191        if cv2.waitKey(1) & 0xFF == ord('q'):
192            break
193            
194    cap.release()
195    cv2.destroyAllWindows()
196
197if __name__ == "__main__":
198    real_time_visual_inference()
199