Agent/examples/find knowledge/scripts/extract_pose_v2.py

#!/usr/bin/env python3
"""
提取人体姿态骨骼图 v2 - 使用MediaPipe Pose
针对不同图片使用不同参数
"""

import mediapipe as mp
import cv2
import numpy as np
import json
import os

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

def extract_pose(image_path, output_dir, img_id, complexity=2, conf=0.5):
    """提取单张图片的姿态骨骼"""
    img = cv2.imread(image_path)
    if img is None:
        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=complexity,
        enable_segmentation=False,
        min_detection_confidence=conf
    ) as pose:
        results = pose.process(img_rgb)
        
        if not results.pose_landmarks:
            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=4, circle_radius=6
            ),
            connection_drawing_spec=mp_drawing.DrawingSpec(
                color=(180, 180, 180), thickness=3
            )
        )
        
        # 保存骨骼图
        skeleton_path = os.path.join(output_dir, f"{img_id}_pose_skeleton.png")
        cv2.imwrite(skeleton_path, skeleton_img)
        
        # 提取关键点坐标
        landmark_names = [lm.name for lm in mp_pose.PoseLandmark]
        landmarks_data = {}
        
        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),
                "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},
                "model_complexity": complexity,
                "detection_confidence": conf,
                "landmarks": landmarks_data,
                "skeleton_image": f"{img_id}_pose_skeleton.png"
            }, f, ensure_ascii=False, indent=2)
        
        return landmarks_data

def main():
    input_dir = "input"
    output_dir = "output/features/pose_skeleton"
    
    # 针对不同图片的参数配置
    # img_5是手臂特写，无法检测全身姿态，记录为局部
    configs = {
        "img_1": {"complexity": 2, "conf": 0.5},
        "img_2": {"complexity": 0, "conf": 0.1},
        "img_3": {"complexity": 2, "conf": 0.5},
        "img_4": {"complexity": 2, "conf": 0.5},
        "img_5": None,  # 手臂特写，无法检测全身
        "img_6": {"complexity": 0, "conf": 0.1},
        "img_7": {"complexity": 0, "conf": 0.1},
        "img_8": {"complexity": 2, "conf": 0.5},
        "img_9": {"complexity": 0, "conf": 0.1},
    }
    
    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):
            continue
        
        config = configs.get(img_id)
        print(f"处理 {img_id}...", end=" ")
        
        if config is None:
            print("跳过（局部特写，无法检测全身姿态）")
            results_summary.append({"image_id": img_id, "detected": False, "reason": "partial_view"})
            continue
        
        landmarks = extract_pose(image_path, output_dir, img_id, 
                                  config["complexity"], config["conf"])
        
        if landmarks:
            print(f"✓ 骨骼图已保存")
            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:
            print("✗ 未检测到姿态")
            results_summary.append({"image_id": img_id, "detected": False, "reason": "not_detected"})
    
    # 建立mapping.json
    pose_segment_map = {
        "img_1": {"segment": "段落1.1", "feature": "女性人物姿态（站立侧身作画）", "element": "元素1"},
        "img_2": {"segment": "段落2.1", "feature": "女性人物姿态（背对镜头作画）", "element": "元素1"},
        "img_3": {"segment": "段落3.1", "feature": "女性人物姿态（跪姿作画）", "element": "元素1"},
        "img_4": {"segment": "段落4.1", "feature": "女性人物姿态（侧身面对画架）", "element": "元素1"},
        "img_6": {"segment": "段落6.1", "feature": "女性人物姿态（背部特写作画）", "element": "元素1"},
        "img_7": {"segment": "段落7.1", "feature": "女性人物姿态（侧颜嗅花）", "element": "元素1"},
        "img_8": {"segment": "段落8.1", "feature": "女性人物姿态（侧身面对画架）", "element": "元素1"},
        "img_9": {"segment": "段落9.1", "feature": "女性人物姿态（背影远景）", "element": "元素1"},
    }
    
    mapping = {
        "dimension": "pose_skeleton",
        "description": "人体姿态骨骼关键点图，使用MediaPipe Pose提取33个关键点",
        "tool": "MediaPipe Pose v0.10.9",
        "format": {
            "skeleton_image": "PNG，黑色背景，白色骨骼连线，尺寸与原图相同",
            "keypoints_json": "JSON，包含33个关键点的归一化坐标(x,y,z)和像素坐标"
        },
        "mappings": []
    }
    
    for result in results_summary:
        img_id = result["image_id"]
        if result["detected"]:
            seg_info = pose_segment_map.get(img_id, {})
            mapping["mappings"].append({
                "file": result["skeleton_file"],
                "keypoints_file": result["keypoints_file"],
                "source_image": f"input/{img_id}.jpg",
                "segment": seg_info.get("segment", ""),
                "category": "实质",
                "feature": seg_info.get("feature", ""),
                "element_id": seg_info.get("element", "元素1"),
                "highlight_cluster": "cluster_1"
            })
        else:
            mapping["mappings"].append({
                "file": None,
                "source_image": f"input/{img_id}.jpg",
                "segment": pose_segment_map.get(img_id, {}).get("segment", ""),
                "category": "实质",
                "feature": "局部特写，无法提取全身姿态",
                "element_id": "元素1",
                "note": result.get("reason", "")
            })
    
    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)
    
    detected = len([r for r in results_summary if r['detected']])
    print(f"\n✓ 完成: {detected}/{len(results_summary)} 张图片检测到姿态")
    print(f"✓ mapping.json 已保存")

if __name__ == "__main__":
    main()