app.py

import os
import cv2
import numpy as np
from collections import defaultdict
import matplotlib.pyplot as plt
from rmn import RMN
import gradio as gr

def process_video(video_path, share_screen_mode):
    # 初始化目录
    output_dir = 'output'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 初始化表情检测模型
    print("Initializing emotion detection model...")
    m = RMN()

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_interval = int(fps * 1)  # 每秒处理一帧
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    print(f"Total frames: {total_frames}, FPS: {fps}")

    # 创建视频写入器
    output_video_path = os.path.join(output_dir, 'output_video.avi')
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    current_frame = 0

    # 面部ID和表情数据
    face_ids = []
    max_face_id = 0
    face_emotions = defaultdict(list)
    max_faces = 0
    initial_faces = []
    last_detections = {}

    print("Starting video processing...")

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Finished processing video.")
            break

        if share_screen_mode:
            # 裁剪右侧1/5区域
            x_start = int(frame_width * 4 / 5)
            frame_to_process = frame[:, x_start:]
        else:
            frame_to_process = frame.copy()
            x_start = 0  # 无偏移

        if current_frame % frame_interval == 0:
            print(f"Processing frame {current_frame}...")
            # 检测面部
            detections = m.detect_faces(frame_to_process)
            print(f"Detected {len(detections)} faces.")
            # 更新最大面部计数
            if len(detections) > max_faces:
                max_faces = len(detections)

            for det in detections:
                xmin = det['xmin']
                ymin = det['ymin']
                xmax = det['xmax']
                ymax = det['ymax']
                matched_id = None
                max_iou = 0

                # 与现有面部进行比较
                for face in initial_faces:
                    ixmin, iymin, ixmax, iymax = face['bbox']
                    # 计算IoU
                    xx1 = max(xmin, ixmin)
                    yy1 = max(ymin, iymin)
                    xx2 = min(xmax, ixmax)
                    yy2 = min(ymax, iymax)
                    inter_area = max(0, xx2 - xx1) * max(0, yy2 - yy1)
                    area1 = (xmax - xmin) * (ymax - ymin)
                    area2 = (ixmax - ixmin) * (iymax - iymin)
                    iou = inter_area / float(area1 + area2 - inter_area + 1e-5)
                    if iou > 0.3 and iou > max_iou:
                        matched_id = face['id']
                        max_iou = iou

                if matched_id is None:
                    if len(initial_faces) < max_faces:
                        # 创建新的面部ID
                        matched_id = max_face_id
                        max_face_id += 1
                        initial_faces.append({'id': matched_id, 'bbox': (xmin, ymin, xmax, ymax)})
                    else:
                        # 基于距离匹配
                        cx = (xmin + xmax) / 2
                        cy = (ymin + ymax) / 2
                        min_dist = float('inf')
                        for face in initial_faces:
                            fx = (face['bbox'][0] + face['bbox'][2]) / 2
                            fy = (face['bbox'][1] + face['bbox'][3]) / 2
                            dist = np.sqrt((cx - fx) ** 2 + (cy - fy) ** 2)
                            if dist < min_dist:
                                min_dist = dist
                                matched_id = face['id']

                # 更新面部边界框
                for face in initial_faces:
                    if face['id'] == matched_id:
                        face['bbox'] = (xmin, ymin, xmax, ymax)
                        break

                # 获取表情标签
                face_img = frame_to_process[ymin:ymax, xmin:xmax]
                if face_img.size == 0:
                    continue
                emo_label, _, _ = m.detect_emotion_for_single_face_image(face_img)
                if emo_label not in ['neutral', 'happy']:
                    emo_label = 'confused'

                # 记录表情
                face_emotions[matched_id].append((current_frame / fps, emo_label))
                print(f"Face {matched_id} emotion: {emo_label}")

                # 更新最后的检测结果，调整坐标到原始帧
                xmin_global = xmin + x_start
                xmax_global = xmax + x_start
                last_detections[matched_id] = (xmin_global, ymin, xmax_global, ymax, emo_label)

        # 在原始帧上绘制最后的检测结果
        for face_id, (xmin, ymin, xmax, ymax, emo_label) in last_detections.items():
            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
            cv2.putText(frame, f"ID:{face_id} {emo_label}", (xmin, ymin + 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

        # 将处理后的帧写入输出视频
        out.write(frame)
        current_frame += 1

    cap.release()
    out.release()

    print("Finished processing video.")

    # 返回输出视频路径和面部表情数据
    return output_video_path, face_emotions

def generate_graphs(selected_ids, face_emotions):
    # 将selected_ids从字符串转换为整数
    selected_ids = [int(face_id) for face_id in selected_ids]
    selected_face_emotions = {face_id: emotions for face_id, emotions in face_emotions.items() if face_id in selected_ids}

    output_dir = 'output'
    emotion_labels = ['confused', 'neutral', 'happy']

    # 生成表情变化图
    plt.figure(figsize=(15, 10))
    for i, (face_id, emotions) in enumerate(selected_face_emotions.items(), 1):
        times = [t for t, _ in emotions]
        labels = [emotion_labels.index(emo) for _, emo in emotions]
        plt.subplot(len(selected_face_emotions), 1, i)
        plt.plot(times, labels, marker='o')
        plt.title(f"Emotion changes for face {face_id}")
        plt.xlabel('Time (s)')
        plt.ylabel('Emotion')
        plt.yticks([0, 1, 2], emotion_labels)
    plt.tight_layout()
    graph_path = os.path.join(output_dir, "selected_faces_emotions.png")
    plt.savefig(graph_path)
    plt.close()
    print("Saved emotion change graph for selected faces.")

    # 生成表情比例图
    time_points = sorted(set(t for emotions in selected_face_emotions.values() for t, _ in emotions))
    emotion_counts_over_time = {t: defaultdict(int) for t in time_points}
    for emotions in selected_face_emotions.values():
        for t, emo in emotions:
            emotion_counts_over_time[t][emo] += 1

    emotion_proportions_over_time = {t: {emo: 0 for emo in emotion_labels} for t in time_points}
    for t in time_points:
        total_faces = sum(emotion_counts_over_time[t].values())
        if total_faces > 0:
            for emo in emotion_labels:
                emotion_proportions_over_time[t][emo] = emotion_counts_over_time[t][emo] / total_faces

    plt.figure(figsize=(15, 10))
    for i, emo in enumerate(emotion_labels, 1):
        proportions = [emotion_proportions_over_time[t][emo] for t in time_points]
        plt.subplot(len(emotion_labels), 1, i)
        plt.plot(time_points, proportions, marker='o')
        plt.title(f"Proportion of {emo} over time")
        plt.xlabel('Time (s)')
        plt.ylabel('Proportion')
        plt.ylim(0, 1)
    plt.tight_layout()
    emotion_proportions_path = os.path.join(output_dir, "selected_emotion_proportions_over_time.png")
    plt.savefig(emotion_proportions_path)
    plt.close()
    print("Saved emotion proportion graph for selected faces.")

    return graph_path, emotion_proportions_path

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Emotion Detection in Videos")

    video_input = gr.Video(label="Upload a video")

    share_screen_checkbox = gr.Checkbox(label="Turn on share mode", value=False)

    process_btn = gr.Button("Process Video")

    video_output = gr.Video(label="Processed Video Output")

    # 状态，用于保存面部表情数据
    face_emotions_state = gr.State()

    # 多选框，列出检测到的ID
    id_checkbox_group = gr.CheckboxGroup(label="Select Face IDs")

    generate_graphs_btn = gr.Button("Generate Graphs")

    graph_output = gr.Image(label="Emotion Change Graph")
    emotion_proportions_output = gr.Image(label="Emotion Proportions Graph")

    def process_and_get_ids(video, share_screen_mode):
        video_output_path, face_emotions = process_video(video, share_screen_mode)
        face_ids = [str(face_id) for face_id in face_emotions.keys()]
        return video_output_path, gr.update(choices=face_ids), face_emotions

    process_btn.click(
        fn=process_and_get_ids,
        inputs=[video_input, share_screen_checkbox],
        outputs=[video_output, id_checkbox_group, face_emotions_state]
    )

    generate_graphs_btn.click(
        fn=generate_graphs,
        inputs=[id_checkbox_group, face_emotions_state],
        outputs=[graph_output, emotion_proportions_output]
    )

demo.launch(share=True)