face_emo_analysize.py

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)