Create face_emo_analysize.py

face_emo_analysize.py

@@ -0,0 +1,283 @@
+import cv2
+import mediapipe as mp
+import math
+import numpy as np
+import time
+import torch
+from PIL import Image
+from torchvision import transforms
+
+
+# 定义预处理函数
+def pth_processing(fp):
+    class PreprocessInput(torch.nn.Module):
+        def __init__(self):
+            super(PreprocessInput, self).__init__()
+
+        def forward(self, x):
+            x = x.to(torch.float32)
+            x = torch.flip(x, dims=(0,))
+            x[0, :, :] -= 91.4953
+            x[1, :, :] -= 103.8827
+            x[2, :, :] -= 131.0912
+            return x
+
+    def get_img_torch(img):
+        ttransform = transforms.Compose([
+            transforms.PILToTensor(),
+            PreprocessInput()
+        ])
+        img = img.resize((224, 224), Image.Resampling.NEAREST)
+        img = ttransform(img)
+        img = torch.unsqueeze(img, 0).to('cuda')
+        return img
+
+    return get_img_torch(fp)
+
+
+# 定义坐标归一化函数
+def norm_coordinates(normalized_x, normalized_y, image_width, image_height):
+    x_px = min(math.floor(normalized_x * image_width), image_width - 1)
+    y_px = min(math.floor(normalized_y * image_height), image_height - 1)
+    return x_px, y_px
+
+
+# 定义获取面部边界框的函数
+def get_box(fl, w, h):
+    idx_to_coors = {}
+    for idx, landmark in enumerate(fl.landmark):
+        landmark_px = norm_coordinates(landmark.x, landmark.y, w, h)
+        if landmark_px:
+            idx_to_coors[idx] = landmark_px
+
+    x_min = np.min(np.asarray(list(idx_to_coors.values()))[:, 0])
+    y_min = np.min(np.asarray(list(idx_to_coors.values()))[:, 1])
+    endX = np.max(np.asarray(list(idx_to_coors.values()))[:, 0])
+    endY = np.max(np.asarray(list(idx_to_coors.values()))[:, 1])
+
+    (startX, startY) = (max(0, x_min), max(0, y_min))
+    (endX, endY) = (min(w - 1, endX), min(h - 1, endY))
+    return startX, startY, endX, endY
+
+
+# 定义显示情感预测结果的函数
+def display_EMO_PRED(img, box, label='', prob=0.0, color=(128, 128, 128), txt_color=(255, 255, 255), line_width=2):
+    lw = line_width or max(round(sum(img.shape) / 2 * 0.003), 2)
+    text2_color = (255, 0, 255)
+    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
+    cv2.rectangle(img, p1, p2, text2_color, thickness=lw, lineType=cv2.LINE_AA)
+    font = cv2.FONT_HERSHEY_SIMPLEX
+
+    tf = max(lw - 1, 1)
+    text_fond = (0, 0, 0)
+
+    # 获取情感标签的文本尺寸
+    label_width, label_height = cv2.getTextSize(label, font, lw / 3, tf)[0]
+
+    # 显示情感标签
+    cv2.putText(img, label,
+                (p1[0], p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,
+                lw / 3, text_fond, thickness=tf, lineType=cv2.LINE_AA)
+    cv2.putText(img, label,
+                (p1[0], p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,
+                lw / 3, text2_color, thickness=tf, lineType=cv2.LINE_AA)
+
+    # 显示情感概率
+    prob_text = f"{prob:.2f}"
+    prob_width, prob_height = cv2.getTextSize(prob_text, font, lw / 3, tf)[0]
+    cv2.putText(img, prob_text,
+                (p1[0] + label_width + 5, p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,
+                lw / 3, text_fond, thickness=tf, lineType=cv2.LINE_AA)
+    cv2.putText(img, prob_text,
+                (p1[0] + label_width + 5, p1[1] - round(((p2[0] - p1[0]) * 20) / 360)), font,
+                lw / 3, text2_color, thickness=tf, lineType=cv2.LINE_AA)
+
+    return img
+
+
+# 定义显示FPS的函数
+def display_FPS(img, text, margin=1.0, box_scale=1.0):
+    img_h, img_w, _ = img.shape
+    line_width = int(min(img_h, img_w) * 0.001)  # line width
+    thickness = max(int(line_width / 3), 1)  # font thickness
+
+    font_face = cv2.FONT_HERSHEY_SIMPLEX
+    font_color = (0, 0, 0)
+    font_scale = thickness / 1.5
+
+    t_w, t_h = cv2.getTextSize(text, font_face, font_scale, None)[0]
+
+    margin_n = int(t_h * margin)
+    sub_img = img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
+              img_w - t_w - margin_n - int(2 * t_h * box_scale): img_w - margin_n]
+
+    white_rect = np.ones(sub_img.shape, dtype=np.uint8) * 255
+
+    img[0 + margin_n: 0 + margin_n + t_h + int(2 * t_h * box_scale),
+    img_w - t_w - margin_n - int(2 * t_h * box_scale):img_w - margin_n] = cv2.addWeighted(sub_img, 0.5, white_rect, .5,
+                                                                                          1.0)
+
+    cv2.putText(img=img,
+                text=text,
+                org=(img_w - t_w - margin_n - int(2 * t_h * box_scale) // 2,
+                     0 + margin_n + t_h + int(2 * t_h * box_scale) // 2),
+                fontFace=font_face,
+                fontScale=font_scale,
+                color=font_color,
+                thickness=thickness,
+                lineType=cv2.LINE_AA,
+                bottomLeftOrigin=False)
+
+    return img
+
+def face_emo_analysize():
+    # 初始化MediaPipe Face Mesh
+    mp_face_mesh = mp.solutions.face_mesh
+
+    # 加载PyTorch模型
+    name = '0_66_49_wo_gl'
+    pth_model = torch.jit.load('torchscript_model_0_66_49_wo_gl.pth'.format(name)).to(
+        'cuda')
+    pth_model.eval()
+
+    # 定义情感字典
+    DICT_EMO = {0: 'Neutral', 1: 'Happiness', 2: 'Sadness', 3: 'Surprise', 4: 'Fear', 5: 'Disgust', 6: 'Anger'}
+
+    # 打开摄像头
+    cap = cv2.VideoCapture(0)
+    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = np.round(cap.get(cv2.CAP_PROP_FPS))
+
+    # 设置视频写入器
+    path_save_video = 'result2.mp4'
+    vid_writer = cv2.VideoWriter(path_save_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+
+    # 使用MediaPipe Face Mesh进行面部检测
+    emotion_stats = {}
+    with mp_face_mesh.FaceMesh(
+            max_num_faces=1,
+            refine_landmarks=False,
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5) as face_mesh:
+        while cap.isOpened():
+            t1 = time.time()
+            success, frame = cap.read()
+            if frame is None: break
+
+            frame_copy = frame.copy()
+            frame_copy.flags.writeable = False
+            frame_copy = cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
+            results = face_mesh.process(frame_copy)
+            frame_copy.flags.writeable = True
+
+            if results.multi_face_landmarks:
+                for fl in results.multi_face_landmarks:
+                    startX, startY, endX, endY = get_box(fl, w, h)
+                    cur_face = frame_copy[startY:endY, startX: endX]
+
+                    # 使用PyTorch模型进行情感预测
+                    cur_face = pth_processing(Image.fromarray(cur_face))
+                    output = torch.nn.functional.softmax(pth_model(cur_face), dim=1).cpu().detach().numpy()[0]
+
+                    # 获取情感类别和概率
+                    cl = np.argmax(output)
+                    label = DICT_EMO[cl]
+                    prob = output[cl]
+
+                    # 记录情感统计信息
+                    if label not in emotion_stats:
+                        emotion_stats[label] = {'start_time': t1, 'duration': 0, 'total_prob': prob, 'count': 1}
+                    else:
+                        emotion_stats[label]['duration'] += (t1 - emotion_stats[label]['start_time'])
+                        emotion_stats[label]['total_prob'] += prob
+                        emotion_stats[label]['count'] += 1
+                        emotion_stats[label]['start_time'] = t1
+
+                    # 显示情感结果和概率
+                    frame = display_EMO_PRED(frame, (startX, startY, endX, endY), label, prob, line_width=3)
+
+            t2 = time.time()
+
+            # 显示FPS
+            frame = display_FPS(frame, 'FPS: {0:.1f}'.format(1 / (t2 - t1)), box_scale=.5)
+
+            # 写入视频
+            vid_writer.write(frame)
+
+            # 显示帧
+            cv2.imshow('Webcam', frame)
+            if cv2.waitKey(1) & 0xFF == ord('\x1b'):
+                break
+
+        # 释放资源
+        vid_writer.release()
+        cap.release()
+        cv2.destroyAllWindows()
+
+        # 打印情感统计信息
+        for emotion, stats in emotion_stats.items():
+            avg_prob = stats['total_prob'] / stats['count']
+            print(f'Emotion: {emotion}, Duration: {stats["duration"]:.2f} seconds, Average Probability: {avg_prob:.2f}')
+
+        # 将视频转换为GIF
+        from moviepy.editor import VideoFileClip
+
+
+        def convert_mp4_to_gif(input_path, output_path, fps=10):
+            clip = VideoFileClip(input_path)
+            clip.write_gif(output_path, fps=fps)
+    #此时我们获得了各表情的持续时间与平均概率,我们可以计算大小,如果负向情绪大于正向情绪那么情感就是负的，再计算平均值即可.
+        positive_emotions = ['Happiness', 'Surprise']
+        negative_emotions = ['Anger', 'Fear', 'Sadness', 'Disgust']
+
+        # 初始化正向和负向情感的统计信息
+        positive_stats = {'duration': 0, 'total_prob': 0, 'count': 0}
+        negative_stats = {'duration': 0, 'total_prob': 0, 'count': 0}
+
+        # 统计正向和负向情感的持续时间和概率
+        for emotion, stats in emotion_stats.items():
+            if emotion in positive_emotions:
+                positive_stats['duration'] += stats['duration']
+                positive_stats['total_prob'] += stats['total_prob']
+                positive_stats['count'] += stats['count']
+            elif emotion in negative_emotions:
+                negative_stats['duration'] += stats['duration']
+                negative_stats['total_prob'] += stats['total_prob']
+                negative_stats['count'] += stats['count']
+
+        # 计算正向和负向情感的平均概率
+        if positive_stats['count'] > 0:
+            positive_avg_prob = positive_stats['total_prob'] / positive_stats['count']
+        else:
+            positive_avg_prob = 0
+
+        if negative_stats['count'] > 0:
+            negative_avg_prob = negative_stats['total_prob'] / negative_stats['count']
+        else:
+            negative_avg_prob = 0
+
+        # 比较正向和负向情感的持续时间
+        if negative_stats['duration'] > positive_stats['duration']:
+            print(f'负向情感持续时间更长: {negative_stats["duration"]:.2f} seconds')
+            print(f'负向情感的平均概率: {negative_avg_prob:.2f}')
+            outcome = "负向,概率:"+str(negative_avg_prob)
+            return outcome
+        else:
+            print(f'正向情感持续时间更长: {positive_stats["duration"]:.2f} seconds')
+            print(f'正向情感的平均概率: {positive_avg_prob:.2f}')
+            outcome = "正向,概率:"+str(positive_avg_prob)
+            return outcome
+        # 将视频转换为GIF
+        from moviepy.editor import VideoFileClip
+
+
+        def convert_mp4_to_gif(input_path, output_path, fps=10):
+            clip = VideoFileClip(input_path)
+            clip.write_gif(output_path, fps=fps)
+
+        # 示例使用
+        input_video_path = "result.mp4"
+        output_gif_path = "result.gif"
+
+        convert_mp4_to_gif(input_video_path, output_gif_path)