fps-arms/source/move.py

import bpy
import math
from random import random
from mathutils import Matrix, Quaternion
from typing import Tuple, Iterable

class RandomBounds:
    def __init__(self, lower_bound, upper_bound):
        self.lower_bound = lower_bound
        self.upper_bound = upper_bound

    def random(self):
        return random() * (self.upper_bound - self.lower_bound) + self.lower_bound

class ArmGraspController:
    """Class to help with controlling a single arm in blender.
    """
    def __init__(self, armature_name: str, wrist_controller_name: str,
                 shoulder_controller_name: str,
                 start_pos: Tuple[float] = (1.6766993999481201, 0.3146645724773407, -1.3483989238739014)):
        """Constructor for ArmGraspController

        Args:
            armature_name (str): The exact name of the blender armature to control.
            wrist_object_name (str): The exact name of the blender object which is used to control the arm's IK solutions.
            start_pos (Tuple[float]): 3D wrist position to start in. (This position is relative to the shoulder.)
        """
        self.armature_name = armature_name
        self.arm = bpy.data.objects[self.armature_name]
        self.open = False
        self.finger_grab_bones = ["index_3", "middle_3", "ring_3", "pinky_3"]
        self.thumb_rot_name = "thumb_3"
        self.upper_arm_name = "upper_arm"
        self.upper_arm_bone = self.arm.pose.bones[self.upper_arm_name]
        self.lower_arm_name = "lower_arm"
        self.lower_arm_bone = self.arm.pose.bones[self.lower_arm_name]
        self.wrist_controller_name = wrist_controller_name
        self.wrist_controller = bpy.data.objects[self.wrist_controller_name]
        self.shoulder_controller_name = shoulder_controller_name
        self.shoulder_controller = bpy.data.objects[self.shoulder_controller_name]
        
        self.optimal_range_factor = 0.95 # Used to modify the max range so that the arm never fully stretches.
        self.arm_range = (self.upper_arm_bone.length + self.lower_arm_bone.length) * self.optimal_range_factor
        
        self.start_pos_wrist = start_pos
        self.start_pos_shoulder = self.get_shoulder_pos().copy()

    def deselect_all_bones(self):
        """Deselects all bones in the armature (the arm).
        """
        for bone in self.arm.pose.bones:    # Deselect all selected bones
            bone.bone.select = False

    @staticmethod
    def vec_add(vec1: Tuple[float], vec2: Tuple[float]):
        return [x+y for (x,y) in zip(vec1, vec2)]

    @staticmethod
    def distance(vec: Iterable[float]):
        return math.sqrt(sum(i**2 for i in vec))

    @staticmethod
    def vec_diff(vec1, vec2):
        return [x-y for (x,y) in zip(vec1, vec2)]

    @staticmethod
    def normalize_vec(vec):
        dist = ArmGraspController.distance(vec)
        return [i/dist for i in vec]

    def get_shoulder_pos(self):
        return self.shoulder_controller.location

    def distance_from_shoulder(self, pos):
        return self.distance(self.vec_diff(pos, self.get_shoulder_pos()))

    def distance_from_shoulder_start(self, pos):
        return self.distance(self.vec_diff(pos, self.start_pos_shoulder))

    def move_wrist(self, new_pos: Tuple[float]):
        """Moves wrist to new position.

        Args:
            new_pos (Tuple[float]): New wrist position, relative to shoulder. Tuple size: 3 floats.
        """
        # self.wrist_object.location = self.vec_add(new_pos, self.upper_arm_bone.head)
        self.wrist_controller.location = new_pos
        self.wrist_controller.keyframe_insert(data_path='location', index=-1)

    def move_shoulder(self, new_pos: Tuple[float]):
        self.shoulder_controller.location = new_pos
        self.shoulder_controller.keyframe_insert(data_path='location', index=-1)

    def set_translation_matrix(self, bone, new_pos):
        for i in range(3):
            bone.matrix[i][3] = new_pos[i]

    def move_arm(self, new_pos: Tuple[float]):
        if(self.distance_from_shoulder(new_pos) <= self.arm_range):
            # Do standard move
            self.move_wrist(new_pos)
        else:
            # Combine both shoulder and arm move
            shoulder_pos = self.get_shoulder_pos()
            diff = self.vec_diff(new_pos, shoulder_pos)
            diff_2 = self.vec_diff(diff, [i*self.arm_range for i in self.normalize_vec(diff)])
            move = self.vec_add(shoulder_pos, diff_2)
            self.move_shoulder(move)

            self.move_wrist(new_pos)

    # Moves the arm, with the shoulder always moving back to its original position.
    def move_arm_rel(self, new_pos: Tuple[float]):
        if(self.distance_from_shoulder_start(new_pos) <= self.arm_range):
            # Do standard move
            self.move_wrist(new_pos)
            # Move shoulder back to start location
            self.move_shoulder(self.start_pos_shoulder)
        else:
            # Combine both shoulder and arm move
            shoulder_pos = self.start_pos_shoulder
            diff = self.vec_diff(new_pos, shoulder_pos)
            diff_2 = self.vec_diff(diff, [i*self.arm_range for i in self.normalize_vec(diff)])
            move = self.vec_add(shoulder_pos, diff_2)
            self.move_shoulder(move)

            self.move_wrist(new_pos)

    def move_back(self):
        self.move_shoulder(self.start_pos_shoulder)
        self.move_wrist(self.start_pos_wrist)

    def open_fingers(self):
        self.grab_movement(0.8)

    def close_fingers(self):
        self.grab_movement(-0.8)

    def move_fingers(self):
        if self.open:
            self.open_fingers()
        else:
            self.close_fingers()
        self.open = not self.open

    def grab_movement(self, angle):
        #  select the 4 base finger bones
        #  Then:
        #  To rotate around local axis (will cause to grab)
        #  Then do the same with thumb, but only around local Z axis!
        bpy.ops.object.mode_set(mode='OBJECT')
        bpy.ops.object.select_all(action='DESELECT') # Deselect all objects
        bpy.context.view_layer.objects.active = self.arm # Make the Armature the active object 
        bpy.ops.object.mode_set(mode='POSE')
        
        # Deselect all fingers
        self.deselect_all_bones()
        
        # Rotate fingers
        for name in self.finger_grab_bones:
            bone = self.arm.pose.bones.get(name)
            # Bones need to be selected when adding keyframes, otherwise blender will throw an error:
            # https://blender.stackexchange.com/questions/1828/what-constitutes-a-context-in-pose-mode
            bone.bone.select = True
            # rotate
            bone.rotation_mode = 'XYZ'
            bone.rotation_euler.rotate_axis('X', angle)
        
        # bpy.ops.transform.rotate(value=angle, orient_axis='X', orient_type='GLOBAL')
        # Add keyframe
        bpy.ops.anim.keying_set_active_set(type='Rotation')
        bpy.ops.anim.keyframe_insert(type='Rotation')
        
        # Now rotate the thumb too
        self.deselect_all_bones()
        bone = self.arm.pose.bones.get(self.thumb_rot_name).bone
        if bone: bone.select = True
        bpy.ops.transform.rotate(value=angle, orient_axis='Z', orient_type='LOCAL')
        bpy.ops.anim.keyframe_insert(type='Rotation')



# Script start
frame_number = 0

max_dist = 0.5
#start_location = ob.location.copy()
start_location = (1.6766993999481201, 0.3146645724773407, -1.3483989238739014)
controller = ArmGraspController("Armature", "WristController", "ShoulderController")
new_pos = [1.7, 2.5, -1.3483989238739014]

x_rand = RandomBounds(0, 2)
y_rand = RandomBounds(1, 5)
z_rand = RandomBounds(-3, 2)

for i in range(10):
    bpy.context.scene.frame_set(frame_number)
    controller.move_arm_rel([x_rand.random(), y_rand.random(), z_rand.random()])
    frame_number += 20