import cv2
import mediapipe as mp
import csv
import os

VIDEO_PATH = r"E:\code\tmp\test.mp4"  # 视频文件路径

# 检查视频文件存在
if not os.path.exists(VIDEO_PATH):
    raise FileNotFoundError(f"视频文件未找到: {VIDEO_PATH}")

# 初始化 MediaPipe Pose 模块
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
pose = mp_pose.Pose(
    static_image_mode=False,
    model_complexity=1,
    enable_segmentation=False,
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5
)

# 打开视频文件
cap = cv2.VideoCapture(VIDEO_PATH)
if not cap.isOpened():
    raise RuntimeError(f"无法打开视频文件: {VIDEO_PATH}")

while True:
    success, frame = cap.read()
    if not success:
        break  # 视频播放完毕

    # 可以调整窗口
    frame = cv2.resize(frame, (1920, 1080))

    # 转为RGB格式
    image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    image_rgb.flags.writeable = False

    # 骨架识别
    results = pose.process(image_rgb)

    # 绘制骨架图像
    image_rgb.flags.writeable = True
    image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    if results.pose_landmarks:
        mp_drawing.draw_landmarks(
            image=image_bgr,
            landmark_list=results.pose_landmarks,
            connections=mp_pose.POSE_CONNECTIONS,
            landmark_drawing_spec=mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
            connection_drawing_spec=mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=2)
        )

    # 显示画面
    cv2.imshow("跳水视频骨架识别（MediaPipe）", image_bgr)
    key = cv2.waitKey(1)
    if key == 27:  # 按Esc键退出
        break

cap.release()
cv2.destroyAllWindows()