Blender_auto_hair_rigid_bodies.py

import bpy
import math
import bmesh
import numpy as np
import sys


# CONSTANT
ARMATURE = "Armature"
HEAD = "Head"
CONTROL_RIG_TAIL = ".001"
RIGID_BODIES_COLLECTION = "Rigid_bodies"
PLAIN_AXESES_COLLECTION = "Plain_axeses"


def check_mode(func_name, mode):
    if bpy.context.mode != mode:
        print(f"{func_name} must be run at {mode} mode ({bpy.context.mode} mode now)")
        return False
    
    return True


def move_collection(obj, collection_name):
    destinationcollection = bpy.data.collections.get(collection_name)

    # Unlink
    for collection in obj.users_collection:
        collection.objects.unlink(obj)

    # Check the collection
    if destinationcollection == None:
        # Making new collection if it is not exist
        destinationcollection = bpy.data.collections.new(collection_name)
        bpy.context.scene.collection.children.link(destinationcollection)
        
    destinationcollection.objects.link(obj)


def determine_rectangular_vertexes(head, long_edge_len, short_edge_len):
    half_short_len = short_edge_len / 2
    return [
        (head[0] - half_short_len, head[2], head[1] + half_short_len),
        (head[0] + half_short_len, head[2], head[1] + half_short_len),
        (head[0] + half_short_len, head[2], head[1] - half_short_len),
        (head[0] - half_short_len, head[2], head[1] - half_short_len),
        (head[0] - half_short_len, head[2] - long_edge_len, head[1] + half_short_len),
        (head[0] + half_short_len, head[2] - long_edge_len, head[1] + half_short_len),
        (head[0] + half_short_len, head[2] - long_edge_len, head[1] - half_short_len),
        (head[0] - half_short_len, head[2] - long_edge_len, head[1] - half_short_len)
    ]


def add_rigid_rectangular(bone):
    # Determination the length of edgies
    head_vec = bone.head[:]
    tail_vec = bone.tail[:]
    bone_vec = np.array([head_vec, tail_vec])
    long_edge_len = np.linalg.norm(bone_vec, ord = 2)
    print(long_edge_len)
    short_edge_len = long_edge_len / 10
    print(long_edge_len, short_edge_len)
            
    # Define the vertexes and faces
    verts = determine_rectangular_vertexes(head_vec, long_edge_len, short_edge_len)
    faces = [(0,1,2,3), (4,5,6,7), (0,4,5,1), (1,5,6,2), (2,6,7,3), (3,7,4,0)]

    # Define the mesh
    mesh = bpy.data.meshes.new("Cube_mesh")
    # Generate the mesh
    mesh.from_pydata(verts,[],faces)
    mesh.update(calc_edges=True)

    # Define the mesh object
    obj = bpy.data.objects.new(f"Rigid_body.{bone.name}", mesh)
    # Link the object into the scene
    bpy.context.scene.collection.objects.link(obj)
    
    return obj


def add_plain_axes(bone):
    obj = bpy.data.objects.new(f"Plain_axes.{bone.name}", None)
    bpy.context.scene.collection.objects.link(obj)
    
    return obj


def add_rigid_bodyies():
    if bpy.context.mode != "POSE":
        print(f"{sys._getframe().f_code.co_name} must be run at pose mode")
        return
    
    rigid_bodies = []
    plain_axeses = []
    bones = bpy.context.selected_pose_bones_from_active_object
    if bones is None:
        print("No bones are found")
        return
    
    for bone in bones:
        print(bone.head, bone.tail)
        rigid_body = add_rigid_rectangular(bone)
        plain_axes = add_plain_axes(bone)
        move_collection(rigid_body, RIGID_BODIES_COLLECTION)
        move_collection(plain_axes, PLAIN_AXESES_COLLECTION)
        rigid_bodies.append(rigid_body)
        plain_axeses.append(plain_axes)
    
    return rigid_bodies, plain_axeses, bones


def attach_to_rigid_bodies(rigid_bodies, plain_axeses, bones):
    if bpy.context.mode != "OBJECT":
        print(f"{sys._getframe().f_code.co_name} must be run at object mode")
        return
    
    for rigid_body, plain_axes, bone in zip(rigid_bodies, plain_axeses, bones):
        # Rigid body
        bpy.ops.object.select_all(action='DESELECT')
        rigid_body.select = True
        bpy.context.view_layer.objects.active = rigid_body
        bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='MEDIAN')
        copy_location = rigid_body.constraints.new(type='COPY_LOCATION')
        copy_location.target = bpy.data.objects["Armature"]
        copy_location.subtarget = bone.name
        copy_location.head_tail = 0.5
        copy_rotation = rigid_body.constraints.new(type='COPY_ROTATION')
        copy_rotation.target = bpy.data.objects["Armature"]
        copy_rotation.subtarget = bone.name
        bpy.ops.rigidbody.object_add()
        if bone.name == HEAD:
            bpy.context.object.rigid_body.type = 'PASSIVE'
            bpy.context.object.rigid_body.kinematic = True
            copy_location.head_tail = 1.0
            continue
                    
        # Plain axes
        bpy.ops.object.select_all(action='DESELECT')
        plain_axes.select = True
        bpy.context.view_layer.objects.active = plain_axes
        bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='MEDIAN')
        copy_location = plain_axes.constraints.new(type='COPY_LOCATION')
        copy_location.target = bpy.data.objects["Armature"]
        copy_location.subtarget = bone.name
        copy_location.head_tail = 0.0
        ## Generic spring
        bpy.ops.rigidbody.constraint_add()
        bpy.context.object.rigid_body_constraint.type = 'GENERIC_SPRING'
        parent_bone = bone.parent
        second_rigid_body_name = f"Rigid_body.{bone.name}"
        first_rigid_body_name = f"Rigid_body.{parent_bone.name}"
        bpy.context.object.rigid_body_constraint.object1 = bpy.data.objects[first_rigid_body_name]
        bpy.context.object.rigid_body_constraint.object2 = bpy.data.objects[second_rigid_body_name]
        bpy.context.object.rigid_body_constraint.use_limit_ang_x = True
        bpy.context.object.rigid_body_constraint.use_limit_ang_y = True
        bpy.context.object.rigid_body_constraint.use_limit_ang_z = True
        bpy.context.object.rigid_body_constraint.use_limit_lin_x = True
        bpy.context.object.rigid_body_constraint.use_limit_lin_y = True
        bpy.context.object.rigid_body_constraint.use_limit_lin_z = True
        bpy.context.object.rigid_body_constraint.use_spring_ang_x = True
        bpy.context.object.rigid_body_constraint.use_spring_ang_y = True
        bpy.context.object.rigid_body_constraint.use_spring_ang_z = True
        bpy.context.object.rigid_body_constraint.use_spring_x = True
        bpy.context.object.rigid_body_constraint.use_spring_y = True
        bpy.context.object.rigid_body_constraint.use_spring_z = True
        bpy.context.object.rigid_body_constraint.limit_ang_z_lower = 0
        bpy.context.object.rigid_body_constraint.limit_ang_z_upper = 0
        bpy.context.object.rigid_body_constraint.limit_ang_x_lower = -0.21
        bpy.context.object.rigid_body_constraint.limit_ang_x_upper = 0.21
        bpy.context.object.rigid_body_constraint.limit_ang_y_lower = 0
        bpy.context.object.rigid_body_constraint.limit_ang_y_upper = 0
        bpy.context.object.rigid_body_constraint.limit_ang_z_lower = -0.436332
        bpy.context.object.rigid_body_constraint.limit_ang_z_upper = 0.436332
        bpy.context.object.rigid_body_constraint.limit_lin_x_upper = 0
        bpy.context.object.rigid_body_constraint.limit_lin_x_lower = 0
        bpy.context.object.rigid_body_constraint.limit_lin_y_upper = 0
        bpy.context.object.rigid_body_constraint.limit_lin_y_lower = 0
        bpy.context.object.rigid_body_constraint.limit_lin_z_upper = 0
        bpy.context.object.rigid_body_constraint.limit_lin_z_lower = 0
        

def find_bone(armature, bone_name):
    if armature.type != 'ARMATURE':
        # Accepting only armature
        return None

    # Get the bone from armature
    target_bone = armature.pose.bones.get(bone_name)
    if target_bone == None:
        # Interruption when no bones are found
        return None

    return target_bone

        
def attach_rigid_to_rig(armature, rigid_bodies, bones):
    mode_safety = check_mode(sys._getframe().f_code.co_name, "POSE")
    if not mode_safety:
        return
    for rigid_body, bone in zip(rigid_bodies, bones):
        if bone.name == HEAD:
            continue
        
        original_bone_name = bone.name.replace(CONTROL_RIG_TAIL, "")
        original_bone = find_bone(armature, original_bone_name)
        if original_bone is None:
            print(f"{original_bone_name} is not found")
            continue
        copy_rotation = original_bone.constraints.new(type='COPY_ROTATION')
        copy_rotation.target = rigid_body
        print(original_bone_name, rigid_body)


# This script must run at pose mode    
if bpy.context.mode == "POSE":
    rigid_bodies, plain_axeses, bones = add_rigid_bodyies()
    bpy.ops.object.mode_set(mode = 'OBJECT')
    attach_to_rigid_bodies(rigid_bodies, plain_axeses, bones)
    armature = bpy.data.objects.get(ARMATURE)
    bpy.context.view_layer.objects.active = armature
    bpy.ops.object.mode_set(mode = 'POSE')
    attach_rigid_to_rig(armature, rigid_bodies, bones)