2024年5月28日火曜日

How to transfer root bone pose animation to the animation of the entire armature

#RootPoseToArmature

import bpy

from mathutils import Matrix, Vector


# Note: This script will automatically set the armature's rotation mode to QUATERNION.

# The root bone will be deleted after the script runs.

# Make sure to back up your project before running this script.


# Set the name of the root bone

root_bone_name = "Root"  # Set the name of the root bone


# Get the currently selected object

armature = bpy.context.active_object

if not armature or armature.type != 'ARMATURE':

    raise ValueError("Please select an armature object.")


bpy.context.view_layer.objects.active = armature


# Switch to Object Mode to get the default pose

bpy.ops.object.mode_set(mode='OBJECT')


# Get the default global location and rotation of the root bone

root_bone = armature.pose.bones.get(root_bone_name)

if not root_bone:

    raise ValueError(f"Root bone '{root_bone_name}' not found.")


# Get the armature's rest pose

armature_matrix_world = armature.matrix_world.copy()

root_bone_matrix_rest = armature.data.bones[root_bone_name].matrix_local.copy()

root_bone_matrix_world_rest = armature_matrix_world @ root_bone_matrix_rest

root_bone_default_loc, root_bone_default_rot, _ = root_bone_matrix_world_rest.decompose()

inverse_root_bone_default_rot = root_bone_default_rot.inverted()


# Get the IK targets

ik_targets = []

for bone in armature.pose.bones:

    for constraint in bone.constraints:

        if constraint.type == 'IK' and constraint.target:

            target = constraint.target

            if target.parent == armature:

                ik_targets.append(target)


# Record the global positions of the IK targets for each keyframe

ik_target_keyframes = {}

for target in ik_targets:

    ik_target_keyframes[target.name] = {}

    if target.animation_data and target.animation_data.action:

        for fcurve in target.animation_data.action.fcurves:

            for keyframe in fcurve.keyframe_points:

                frame = int(keyframe.co.x)

                bpy.context.scene.frame_set(frame)

                ik_target_keyframes[target.name][frame] = target.matrix_world.translation.copy()


# Get the animation data

animation_data = armature.animation_data

if not animation_data:

    raise ValueError("Animation data not found.")


action = animation_data.action

if not action:

    raise ValueError("Action not found in animation data.")


# Process the animation data and apply the root bone transform to the armature

transform_data = {'location': [], 'rotation_quaternion': [], 'scale': []}

for fcurve in action.fcurves:

    if fcurve.data_path.startswith(f"pose.bones[\"{root_bone_name}\"]"):

        for keyframe in fcurve.keyframe_points:

            frame = int(keyframe.co.x)


            # Move to the frame to set the keyframe

            bpy.context.scene.frame_set(frame)


            # Get the pose position of the root bone

            root_bone_matrix = root_bone.matrix


            # Calculate the global transform of the root bone

            global_matrix = armature.matrix_world @ root_bone_matrix


            # Decompose the matrix into location, rotation, and scale

            global_loc, global_rot, global_scale = global_matrix.decompose()


            # Save the transform data in the dictionary

            transform_data['location'].append((frame, global_loc))

            transform_data['rotation_quaternion'].append((frame, global_rot))

            transform_data['scale'].append((frame, global_scale))


# Switch to Object Mode to apply transformations

bpy.ops.object.mode_set(mode='OBJECT')


# Automatically set the armature's rotation mode to QUATERNION

armature.rotation_mode = 'QUATERNION'


# Apply the transforms to the armature object and insert keyframes

for (frame, loc), (_, rot) in zip(transform_data['location'], transform_data['rotation_quaternion']):

    # Correct the rotation for the root bone's default rotation

    corrected_rot = (rot @ inverse_root_bone_default_rot)

    

    # Create the translation matrices for the default and current locations

    translation_matrix_to_default = Matrix.Translation(-root_bone_default_loc)

    translation_matrix_back = Matrix.Translation(root_bone_default_loc)

    translation_matrix_current = Matrix.Translation(loc)

    

    # Create the rotation matrix for the corrected rotation

    rotation_matrix = corrected_rot.to_matrix().to_4x4()

    

    # Combine the matrices: 

    # 1. Translate to default position

    # 2. Apply rotation

    # 3. Translate to the current location

    final_matrix = translation_matrix_current @ rotation_matrix @ translation_matrix_to_default

    

    # Decompose the final matrix into location and rotation

    final_loc, final_rot, _ = final_matrix.decompose()

    

    # Apply the final location and rotation to the armature

    armature.location = final_loc

    armature.rotation_quaternion = final_rot

    armature.keyframe_insert(data_path="location", frame=frame)

    armature.keyframe_insert(data_path="rotation_quaternion", frame=frame)


for frame, scale in transform_data['scale']:

    armature.scale = scale

    armature.keyframe_insert(data_path="scale", frame=frame)


# Update the IK target positions for each frame

for target in ik_targets:

    for frame, initial_position in ik_target_keyframes[target.name].items():

        bpy.context.scene.frame_set(frame)

        # Convert the initial position to the armature's local space at the current frame

        local_pos = armature.matrix_world.inverted() @ initial_position

        target.location = local_pos

        target.keyframe_insert(data_path="location", frame=frame)


# Delete the root bone

bpy.ops.object.mode_set(mode='EDIT')

bones = armature.data.edit_bones

root_bone = bones.get(root_bone_name)

if root_bone:

    bones.remove(root_bone)


# Restore the current frame

bpy.context.scene.frame_set(bpy.context.scene.frame_current)