109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
import math
|
|
from hsemotion_onnx.facial_emotions import HSEmotionRecognizer
|
|
|
|
EMOTION_VA_MAP = {
|
|
'happy': (0.85, 0.60),
|
|
'sad': (-0.75, -0.60),
|
|
'angry': (-0.70, 0.80),
|
|
'fear': (-0.65, 0.75),
|
|
'surprise': (0.20, 0.85),
|
|
'disgust': (-0.80, 0.40),
|
|
'neutral': (0.00, 0.00),
|
|
'contempt': (-0.60, 0.50),
|
|
}
|
|
|
|
EMOTION_HANDLE = {
|
|
'happiness': 'happy',
|
|
'sadness': 'sad',
|
|
'anger': 'angry',
|
|
'fear': 'fear',
|
|
'surprise': 'surprise',
|
|
'disgust': 'disgust',
|
|
'neutral': 'neutral',
|
|
'contempt': 'contempt',
|
|
}
|
|
|
|
def get_fine_grained_emotion(valence, arousal):
|
|
radius = math.sqrt(valence**2 + arousal**2)
|
|
angle = math.degrees(math.atan2(arousal, valence))
|
|
|
|
if radius < 0.25:
|
|
return "neutral"
|
|
|
|
if 0 <= angle < 90:
|
|
if angle > 60: return "excited"
|
|
elif angle > 30: return "happy"
|
|
else: return "pleased"
|
|
elif 90 <= angle <= 180:
|
|
if angle > 150: return "nervous"
|
|
elif angle > 120: return "angry"
|
|
else: return "annoying"
|
|
elif -180 <= angle < -90:
|
|
if angle < -150: return "sad"
|
|
elif angle < -120: return "bored"
|
|
else: return "sleepy"
|
|
elif -90 <= angle < 0:
|
|
if angle < -60: return "calm"
|
|
elif angle < -30: return "peaceful"
|
|
else: return "relaxed"
|
|
|
|
return "neutral"
|
|
|
|
class EmotionAnalyzer:
|
|
def __init__(self):
|
|
print("正在加载 HSEmotion-ONNX 模型...")
|
|
self.fer = HSEmotionRecognizer(model_name='enet_b0_8_best_vgaf')
|
|
print("HSEmotion-ONNX 模型加载完成")
|
|
|
|
def calculate_va_score(self, emotion_prob):
|
|
valence_sum = 0.0
|
|
arousal_sum = 0.0
|
|
total_prob = 0.0
|
|
|
|
for emotion, prob in emotion_prob.items():
|
|
key = EMOTION_HANDLE.get(emotion.lower(), emotion.lower())
|
|
if key in EMOTION_VA_MAP:
|
|
v, a = EMOTION_VA_MAP[key]
|
|
valence_sum += v * prob
|
|
arousal_sum += a * prob
|
|
total_prob += prob
|
|
|
|
if total_prob == 0:
|
|
return 0.0, 0.0
|
|
return valence_sum, arousal_sum
|
|
|
|
def analyze(self, face_img_bgr):
|
|
if face_img_bgr is None or face_img_bgr.size == 0:
|
|
return []
|
|
|
|
face_img_rgb = cv2.cvtColor(face_img_bgr, cv2.COLOR_BGR2RGB)
|
|
|
|
# predict_emotions 返回主要情绪标签和概率数组
|
|
emotion_raw, scores = self.fer.predict_emotions(face_img_rgb, logits=False)
|
|
|
|
probabilities = {}
|
|
for idx, score in enumerate(scores):
|
|
raw_label = self.fer.idx_to_class[idx]
|
|
key = EMOTION_HANDLE.get(raw_label.lower(), raw_label.lower())
|
|
|
|
probabilities[key] = float(score)
|
|
|
|
valence, arousal = self.calculate_va_score(probabilities)
|
|
fine_grained_label = get_fine_grained_emotion(valence, arousal)
|
|
|
|
result = {
|
|
"box": {},
|
|
"vaVal": (round(valence, 4), round(arousal, 4)),
|
|
"probabilities": probabilities,
|
|
"dominant_emotion": EMOTION_HANDLE.get(emotion_raw.lower(), emotion_raw.lower()),
|
|
"emotion": fine_grained_label
|
|
}
|
|
|
|
return [result]
|
|
|
|
analyzer_instance = EmotionAnalyzer()
|
|
|
|
def analyze_emotion_with_hsemotion(face_crop_bgr):
|
|
return analyzer_instance.analyze(face_crop_bgr) |