Spaces:
Sleeping
Sleeping
File size: 3,066 Bytes
1e91e1d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import numpy as np
class SkeletonGenerator:
def __init__(self, custom_mapping=None):
# Define default Maya-compatible bone names and their corresponding MediaPipe indices
self.bone_mapping = {
# Root & Spine
'Hips': 23, # HIPS/Pelvis
'Spine': 24, # Lower spine
'Spine1': 12, # Mid spine
'Spine2': 11, # Upper spine
# Head & Neck
'Neck': 11, # BASE_NECK
'Head': 0, # HEAD (using nose as reference)
# Left Arm Chain
'LeftShoulder': 11, # LEFT_SHOULDER
'LeftArm': 13, # LEFT_UPPER_ARM
'LeftForeArm': 15, # LEFT_LOWER_ARM
'LeftHand': 15, # LEFT_HAND (using wrist as reference)
# Right Arm Chain
'RightShoulder': 12, # RIGHT_SHOULDER
'RightArm': 14, # RIGHT_UPPER_ARM
'RightForeArm': 16, # RIGHT_LOWER_ARM
'RightHand': 16, # RIGHT_HAND (using wrist as reference)
# Left Leg Chain
'LeftUpLeg': 23, # LEFT_UPPER_LEG
'LeftLeg': 25, # LEFT_LOWER_LEG
'LeftFoot': 27, # LEFT_FOOT
'LeftToeBase': 31, # LEFT_TOE
# Right Leg Chain
'RightUpLeg': 24, # RIGHT_UPPER_LEG
'RightLeg': 26, # RIGHT_LOWER_LEG
'RightFoot': 28, # RIGHT_FOOT
'RightToeBase': 32, # RIGHT_TOE
}
def generate_skeleton(self, landmarks):
"""
Convert MediaPipe landmarks to Maya-compatible skeleton data
"""
if landmarks is None:
return None
skeleton_data = {}
# Convert landmarks to skeleton joint positions
for bone_name, landmark_idx in self.bone_mapping.items():
position = landmarks[landmark_idx]
rotation = self._calculate_bone_rotation(landmarks, landmark_idx)
# Convert NumPy arrays to regular Python lists
skeleton_data[bone_name] = {
'position': position.tolist(),
'rotation': rotation.tolist()
}
return skeleton_data
def _calculate_bone_rotation(self, landmarks, landmark_idx):
"""
Calculate bone rotation based on connected landmarks
"""
# Simple rotation calculation - can be enhanced for better accuracy
rotation = np.zeros(3)
if landmark_idx > 0:
# Calculate rotation based on parent-child relationship
parent_idx = landmark_idx - 1
direction = landmarks[landmark_idx] - landmarks[parent_idx]
# Convert direction to euler angles (simplified)
rotation[0] = np.arctan2(direction[1], direction[2]) # pitch
rotation[1] = np.arctan2(direction[0], direction[2]) # yaw
rotation[2] = np.arctan2(direction[0], direction[1]) # roll
return rotation |