#!/usr/bin/env python3
"""
提取人体姿态骨骼图 - 使用MediaPipe Pose
输出：每张图片的骨骼关键点图（PNG）+ 关键点坐标（JSON）
"""

import mediapipe as mp
import cv2
import numpy as np
import json
import os
from PIL import Image

mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

def extract_pose(image_path, output_dir, img_id):
    """提取单张图片的姿态骨骼"""
    img = cv2.imread(image_path)
    if img is None:
        print(f"无法读取图片: {image_path}")
        return None
    
    h, w = img.shape[:2]
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    with mp_pose.Pose(
        static_image_mode=True,
        model_complexity=2,
        enable_segmentation=False,
        min_detection_confidence=0.5
    ) as pose:
        results = pose.process(img_rgb)
        
        if not results.pose_landmarks:
            print(f"  未检测到姿态: {img_id}")
            return None
        
        # 创建黑色背景的骨骼图
        skeleton_img = np.zeros((h, w, 3), dtype=np.uint8)
        
        # 绘制骨骼连接线（白色）
        mp_drawing.draw_landmarks(
            skeleton_img,
            results.pose_landmarks,
            mp_pose.POSE_CONNECTIONS,
            landmark_drawing_spec=mp_drawing.DrawingSpec(
                color=(255, 255, 255), thickness=3, circle_radius=5
            ),
            connection_drawing_spec=mp_drawing.DrawingSpec(
                color=(200, 200, 200), thickness=2
            )
        )
        
        # 保存骨骼图
        skeleton_path = os.path.join(output_dir, f"{img_id}_pose_skeleton.png")
        cv2.imwrite(skeleton_path, skeleton_img)
        
        # 提取关键点坐标
        landmarks_data = {}
        landmark_names = [lm.name for lm in mp_pose.PoseLandmark]
        
        for i, landmark in enumerate(results.pose_landmarks.landmark):
            name = landmark_names[i] if i < len(landmark_names) else f"landmark_{i}"
            landmarks_data[name] = {
                "x": round(landmark.x, 4),  # 归一化坐标 [0,1]
                "y": round(landmark.y, 4),
                "z": round(landmark.z, 4),
                "visibility": round(landmark.visibility, 4),
                "pixel_x": int(landmark.x * w),
                "pixel_y": int(landmark.y * h)
            }
        
        # 保存关键点JSON
        json_path = os.path.join(output_dir, f"{img_id}_pose_keypoints.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump({
                "image_id": img_id,
                "image_size": {"width": w, "height": h},
                "landmarks": landmarks_data,
                "skeleton_image": f"{img_id}_pose_skeleton.png"
            }, f, ensure_ascii=False, indent=2)
        
        print(f"  ✓ {img_id}: 骨骼图已保存 -> {skeleton_path}")
        return landmarks_data

def main():
    input_dir = "input"
    output_dir = "output/features/pose_skeleton"
    
    results_summary = []
    
    for i in range(1, 10):
        img_id = f"img_{i}"
        image_path = os.path.join(input_dir, f"{img_id}.jpg")
        
        if not os.path.exists(image_path):
            print(f"图片不存在: {image_path}")
            continue
        
        print(f"处理 {img_id}...")
        landmarks = extract_pose(image_path, output_dir, img_id)
        
        if landmarks:
            results_summary.append({
                "image_id": img_id,
                "detected": True,
                "keypoints_file": f"{img_id}_pose_keypoints.json",
                "skeleton_file": f"{img_id}_pose_skeleton.png"
            })
        else:
            results_summary.append({
                "image_id": img_id,
                "detected": False
            })
    
    # 保存mapping.json
    mapping = {
        "dimension": "pose_skeleton",
        "description": "人体姿态骨骼关键点图，使用MediaPipe Pose提取33个关键点",
        "tool": "MediaPipe Pose v0.10.9",
        "format": {
            "skeleton_image": "PNG，黑色背景，白色骨骼连线",
            "keypoints_json": "JSON，包含33个关键点的归一化坐标和像素坐标"
        },
        "mappings": []
    }
    
    # 根据制作表结构建立对应关系
    pose_segment_map = {
        "img_1": [
            {"segment": "段落1.1", "category": "实质", "feature": "女性人物姿态", "element": "元素1"},
        ],
        "img_2": [
            {"segment": "段落2.1", "category": "实质", "feature": "女性人物姿态", "element": "元素1"},
        ],
        "img_3": [
            {"segment": "段落3.1", "category": "实质", "feature": "女性人物姿态（跪姿）", "element": "元素1"},
        ],
        "img_4": [
            {"segment": "段落4.1", "category": "实质", "feature": "女性人物姿态（侧身）", "element": "元素1"},
        ],
        "img_5": [
            {"segment": "段落5.1", "category": "实质", "feature": "女性人物姿态（手臂特写）", "element": "元素1"},
        ],
        "img_6": [
            {"segment": "段落6.1", "category": "实质", "feature": "女性人物姿态（背部特写）", "element": "元素1"},
        ],
        "img_7": [
            {"segment": "段落7.1", "category": "实质", "feature": "女性人物姿态（侧颜/嗅花）", "element": "元素1"},
        ],
        "img_8": [
            {"segment": "段落8.1", "category": "实质", "feature": "女性人物姿态（侧身）", "element": "元素1"},
        ],
        "img_9": [
            {"segment": "段落9.1", "category": "实质", "feature": "女性人物姿态（背影远景）", "element": "元素1"},
        ],
    }
    
    for result in results_summary:
        img_id = result["image_id"]
        if result["detected"]:
            segments = pose_segment_map.get(img_id, [])
            for seg in segments:
                mapping["mappings"].append({
                    "file": result["skeleton_file"],
                    "keypoints_file": result["keypoints_file"],
                    "source_image": f"input/{img_id}.jpg",
                    "segment": seg["segment"],
                    "category": seg["category"],
                    "feature": seg["feature"],
                    "element_id": seg["element"]
                })
    
    mapping_path = os.path.join(output_dir, "mapping.json")
    with open(mapping_path, 'w', encoding='utf-8') as f:
        json.dump(mapping, f, ensure_ascii=False, indent=2)
    
    print(f"\n✓ mapping.json 已保存: {mapping_path}")
    print(f"✓ 处理完成: {len([r for r in results_summary if r['detected']])} / {len(results_summary)} 张图片检测到姿态")

if __name__ == "__main__":
    main()