#!/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()