from typing import Union, List, Optional
import numpy as np
from mathutils import Vector, Euler, Color, Matrix, Quaternion
import bpy
from blenderproc.python.utility.Utility import KeyFrame
from blenderproc.python.types.EntityUtility import Entity
from blenderproc.python.types.MeshObjectUtility import MeshObject
from blenderproc.python.types.BoneUtility import get_constraint, set_ik_constraint, \
set_ik_limits_from_rotation_constraint
from blenderproc.python.types.InertialUtility import Inertial
[docs]class Link(Entity):
def __init__(self, bpy_object: bpy.types.Object):
super().__init__(bpy_object)
object.__setattr__(self, 'visuals', [])
object.__setattr__(self, 'inertial', None)
object.__setattr__(self, 'collisions', [])
object.__setattr__(self, 'joint_type', None)
object.__setattr__(self, 'bone', None)
object.__setattr__(self, 'fk_bone', None)
object.__setattr__(self, 'ik_bone', None)
object.__setattr__(self, 'ik_bone_controller', None)
object.__setattr__(self, 'ik_bone_constraint', None)
object.__setattr__(self, 'armature', None)
object.__setattr__(self, 'child', None)
object.__setattr__(self, 'fk_ik_mode', "fk")
object.__setattr__(self, 'link2bone_mat', None)
object.__setattr__(self, 'visual_local2link_mats', [])
object.__setattr__(self, 'collision_local2link_mats', [])
object.__setattr__(self, 'inertial_local2link_mat', None)
[docs] def set_parent(self, parent: "Link"):
""" Sets the parent of this link.
:param parent: Parent link.
"""
assert isinstance(parent, Link)
object.__setattr__(self, "parent", parent)
[docs] def get_parent(self) -> "Link":
""" Returns this link's parent.
:return: Parent link.
"""
return self.parent
[docs] def set_child(self, child: "Link"):
""" Sets the child of this link.
:param child: Child link.
"""
assert isinstance(child, Link)
object.__setattr__(self, "child", child)
[docs] def get_child(self) -> "Link":
""" Returns this link's child.
:return: Child link.
"""
return self.child
[docs] def _set_rotation_euler(self, bone: bpy.types.PoseBone,
rotation_euler: Union[float, List[float], Euler, np.ndarray], mode: str = "absolute"):
""" Rotates the bone based on euler angles. Validates values with given constraints.
:param bone: The bone to be rotated.
:param rotation_euler: The amount of rotation (in radians). Either three floats for x, y and z axes, or a
single float. In the latter case, the axis of rotation is derived based on the rotation
constraint. If these are not properly set (i.e., two axes must have equal min/max
values) an exception will be thrown.
:param mode: One of ["absolute", "relative"]. For absolute rotations we clip the rotation value based on the
constraints. For relative we don't - this will result in inverse motion after the constraint's
limits have been reached.
"""
assert mode in ["absolute", "relative"]
bpy.ops.object.select_all(action='DESELECT')
bone.bone.select = True
bone.rotation_mode = 'XYZ'
# in absolute mode we overwrite the rotation values of the armature
if mode == "absolute":
if isinstance(rotation_euler, float):
axis = self._determine_rotation_axis(bone=bone)
rotation_euler = self._clip_value_from_constraint(bone=bone, value=rotation_euler,
constraint_name="Limit Rotation", axis=axis)
current_rotation_euler = bone.rotation_euler
current_rotation_euler[["X", "Y", "Z"].index(axis)] = rotation_euler
bone.rotation_euler = current_rotation_euler
print(f"Set rotation_euler of bone {bone.name} to {rotation_euler}")
else:
bone.rotation_euler = Vector(
[self._clip_value_from_constraint(bone=bone, value=rot_euler, constraint_name="Limit Rotation",
axis=axis)
for rot_euler, axis in zip(rotation_euler, ["X", "Y", "Z"])])
print(f"Set rotation_euler of bone {bone.name} to {rotation_euler}")
# in relative mode we add the rotation to the current value
elif mode == "relative":
if isinstance(rotation_euler, float):
axis = self._determine_rotation_axis(bone=bone)
if axis is not None:
bone.rotation_euler.rotate_axis(axis, rotation_euler)
else:
for axis in ['X', 'Y', 'Z']:
bone.rotation_euler.rotate_axis(axis, rotation_euler)
print(f"Relatively rotated bone {bone.name} around axis {axis} for {rotation_euler} radians")
else:
for axis, rotation in zip(["X", "Y", "Z"], rotation_euler):
bone.rotation_euler.rotate_axis(axis, rotation)
print(f"Relatively rotated {bone.name} for {rotation_euler} radians")
[docs] def set_rotation_euler(self, *args, **kwargs):
raise NotImplementedError("Please use 'set_rotation_euler_fk()' or set_rotation_euler_ik()'!")
[docs] def set_location(self, *args, **kwargs):
raise NotImplementedError("Please use 'set_location_ik()'!")
[docs] def set_rotation_euler_fk(self, *args, **kwargs):
""" Sets the rotation for this link in forward kinematics mode. See self._set_rotation_euler() for details. """
if self.get_fk_ik_mode() != "fk":
self.switch_fk_ik_mode(mode="fk")
self._set_rotation_euler(bone=self.fk_bone, *args, **kwargs)
[docs] def set_rotation_euler_ik(self, *args, **kwargs):
""" Sets the rotation for this link in inverse kinematics mode. See self._set_rotation_euler() for details. """
if self.get_fk_ik_mode() != "ik":
self.switch_fk_ik_mode(mode="ik")
self._set_rotation_euler(bone=self.ik_bone_controller, *args, **kwargs)
[docs] def set_location_ik(self, location: Union[List[float], np.array, Vector]):
""" Sets the location of the ik bone controller in inverse kinematics mode.
:param location: Location vector.
"""
if self.get_fk_ik_mode() != "ik":
self.switch_fk_ik_mode(mode="ik")
if not isinstance(location, Vector):
location = Vector(location)
assert self.ik_bone_controller is not None, f"No ik bone chain created. Please run " \
f"'urdf_object.create_ik_bone_controller()' first!"
# first: determine offset to base frame
self.ik_bone_controller.location = Vector([0., 0., 0.])
self.ik_bone_controller.rotation_mode = "XYZ"
self.ik_bone_controller.rotation_euler = Vector([0., 0., 0.])
bpy.context.view_layer.update()
offset_mat = self.ik_bone_controller.matrix
# offset location by ik offset to base bones
location = location + (self.ik_bone.head - self.bone.head)
location = Vector([location[0], location[1], location[2], 1.])
# move to current frame
location = offset_mat.inverted() @ location
self.ik_bone_controller.location = location[:3]
[docs] def _determine_rotation_axis(self, bone: bpy.types.PoseBone) -> Optional[str]:
""" Determines the single rotation axis and checks if the constraints are set well to have only one axis of
freedom.
:param bone: Bone of which the rotation axis will be determined.
:return: The single rotation axis ('X', 'Y' or 'Z') or None if no constraint is set..
"""
c = get_constraint(bone=bone, constraint_name="Limit Rotation")
if c is None:
print("WARNING: No rotation constraint set. Will rotate all axes relatively!")
return None
axes = ['X', 'Y', 'Z']
if c.use_limit_x and c.min_x == c.max_x:
axes.pop(axes.index('X'))
if c.use_limit_y and c.min_y == c.max_y:
axes.pop(axes.index('Y'))
if c.use_limit_z and c.min_z == c.max_z:
axes.pop(axes.index('Z'))
assert len(axes) == 1, f"Constraints are set wrong for a rotation around a single axis. Only one axis should " \
f"be allowed to move, but found freedom in {len(axes)} axes of armature " \
f"{self.get_name()} (constraint: {c}, uses limits (xyz): " \
f"{c.use_limit_x, c.use_limit_y, c.use_limit_z}, " \
f"values: {c.min_x, c.max_x, c.min_y, c.max_y, c.min_z, c.max_z})."
return axes[0]
[docs] def _clip_value_from_constraint(self, bone: bpy.types.PoseBone, value: float = 0, constraint_name: str = "",
axis: str = "X") -> float:
""" Checks if an axis is constraint, and clips the value to the min/max of this constraint. If the constraint
does not exist, nothing is done.
:param bone: The bone from which the constraints will be determined.
:param value: Value to be clipped.
:param constraint_name: Name of the constraint.
:param axis: Axis to check.
:return: Clipped value if a constraint is set, else the initial value.
"""
c = get_constraint(bone=bone, constraint_name=constraint_name)
if c is not None:
min_value = eval(f"c.min_{axis.lower()}")
max_value = eval(f"c.max_{axis.lower()}")
print(f"Clipping {value} to be in range {min_value}, {max_value}")
if value < min_value:
return min_value
elif value > max_value:
return max_value
return value
[docs] def set_visuals(self, visuals: List[MeshObject]):
""" Sets the visual meshes for this link.
:param visuals: List of visual meshes.
"""
object.__setattr__(self, "visuals", visuals)
[docs] def get_visuals(self) -> List[MeshObject]:
""" Returns the visual meshes for this link.
:return: List of visual meshes.
"""
return self.visuals
[docs] def set_collisions(self, collisions: List[MeshObject]):
""" Sets the collision meshes for this link.
:param collisions: List of collision meshes.
"""
object.__setattr__(self, "collisions", collisions)
[docs] def get_collisions(self) -> List[MeshObject]:
""" Returns the collision meshes for this link.
:return: List of collision meshes.
"""
return self.collisions
[docs] def set_inertial(self, inertial: Inertial):
""" Sets the inertial meshes for this link.
:param inertial: List of inertial meshes.
"""
object.__setattr__(self, "inertial", inertial)
[docs] def get_inertial(self) -> Inertial:
""" Returns the inertial meshes for this link.
:return: List of inertial meshes.
"""
return self.inertial
[docs] def set_bone(self, bone: bpy.types.PoseBone):
""" Sets the bone controlling the visuals / collisions / inertial of the link.
:param bone: The bone.
"""
object.__setattr__(self, "bone", bone)
[docs] def set_fk_bone(self, bone: bpy.types.PoseBone):
""" Sets the bone controlling the forward kinematic motion of this link.
:param bone: The bone.
"""
object.__setattr__(self, "fk_bone", bone)
[docs] def set_ik_bone(self, bone: bpy.types.PoseBone):
""" Sets the bone controlling the inverse kinematic motion of this link.
:param bone: The bone.
"""
object.__setattr__(self, "ik_bone", bone)
[docs] def set_ik_bone_controller(self, bone: bpy.types.PoseBone):
""" Sets the control bone controlling the inverse kinematic motion for this link.
:param bone: The bone.
"""
object.__setattr__(self, "ik_bone_controller", bone)
[docs] def set_ik_bone_constraint(self, bone: bpy.types.PoseBone):
""" Sets the constraint bone responsible for constraining to inverse kinematic motion.
:param bone: The bone.
"""
object.__setattr__(self, "ik_bone_constraint", bone)
[docs] def set_armature(self, armature: bpy.types.Armature):
""" Sets the armature which holds all the bones of all links.
:param armature: The armature.
"""
object.__setattr__(self, "armature", armature)
[docs] def _set_fk_ik_mode(self, mode="fk"):
""" Sets the mode of the link.
:param mode: One of ["fk", "ik"] denoting forward or inverse kinematics mode.
"""
object.__setattr__(self, "fk_ik_mode", mode)
[docs] def set_link2bone_mat(self, matrix: Matrix):
""" Sets the transformation matrix from bone to link.
:param matrix: The transformation from bone to link.
"""
object.__setattr__(self, "link2bone_mat", matrix)
[docs] def set_visual_local2link_mats(self, matrix_list: List[Matrix]):
""" Sets the transformation matrices from link to the visual parts.
:param matrix_list: List of transformation matrices.
"""
object.__setattr__(self, "visual_local2link_mats", matrix_list)
[docs] def set_collision_local2link_mats(self, matrix_list: List[Matrix]):
""" Sets the transformation matrices from link to the collision parts.
:param matrix_list: List of transformation matrices.
"""
object.__setattr__(self, "collision_local2link_mats", matrix_list)
[docs] def set_inertial_local2link_mat(self, matrix: Matrix):
""" Sets the transformation matrix from link to inertial.
:param matrix: The transformation matrix from link to inertial.
"""
object.__setattr__(self, "inertial_local2link_mat", matrix)
[docs] def hide(self, hide_object: bool = True):
""" Sets the visibility of the object and all visual, collision and inertial parts.
:param hide_object: Determines whether the object should be hidden in rendering.
"""
self.hide(hide_object=hide_object)
for obj in self.get_all_objs():
obj.hide(hide_object=hide_object)
[docs] def get_visual_local2world_mats(self) -> Optional[List]:
"""Returns the transformation matrices from world to the visual parts.
:return: List of transformation matrices.
"""
bpy.context.view_layer.update()
bone_mat = Matrix.Identity(4)
if self.bone is not None:
bone_mat = self.bone.matrix
if not self.visuals:
return [bone_mat @ (self.link2bone_mat.inverted() @ visual_local2link_mat) for visual_local2link_mat in
self.visual_local2link_mats]
else:
return None
[docs] def get_collision_local2world_mats(self) -> Optional[List]:
"""Returns the transformation matrices from world to the collision parts.
:return: List of transformation matrices.
"""
bpy.context.view_layer.update()
bone_mat = Matrix.Identity(4)
if self.bone is not None:
bone_mat = self.bone.matrix
if not self.collisions:
return [bone_mat @ (self.link2bone_mat.inverted() @ collision_local2link_mat) for collision_local2link_mat
in self.collision_local2link_mats]
else:
return None
[docs] def get_inertial_local2world_mat(self) -> Matrix:
"""Returns the transformation matrix from world to the inertial part.
:return: The transformation matrix.
"""
bpy.context.view_layer.update()
bone_mat = Matrix.Identity(4)
if self.bone is not None:
bone_mat = self.bone.matrix
if self.inertial is not None:
return bone_mat @ (self.link2bone_mat.inverted() @ self.inertial_local2link_mat)
else:
return None
[docs] def get_all_objs(self):
""" Returns all meshes of this link.
:return: List of all meshes (visual, collision, inertial) for this link.
"""
return self.visuals + self.collisions + ([self.inertial] if self.inertial is not None else [])
[docs] def set_joint_type(self, joint_type: Optional[str]):
""" Sets the joint type of the link which specifies the connection to its parent.
:param joint_type: One of ['fixed', 'prismatic', 'revolute', 'continuous', 'planar', 'floating' or None].
"""
object.__setattr__(self, "joint_type", joint_type)
[docs] def get_joint_type(self) -> Optional[str]:
""" Returns the joint type.
:return: The joint type of the armature.
"""
return self.joint_type
[docs] def parent_with_bone(self, weight_distribution='rigid'):
""" Parents all objects of the link to the bone.
:param weight_distribution: One of ['envelope', 'automatic', 'rigid']. For more information please see
https://docs.blender.org/manual/en/latest/animation/armatures/skinning/parenting.html.
"""
assert weight_distribution in ['envelope', 'automatic', 'rigid']
if self.bone is None:
# only set parent
print(f"WARNING: Link {self.get_name()} does not have a bone to be parented with (usually because it's the "
f"first link)!")
self.blender_obj.parent = self.armature
return
if weight_distribution in ['envelope', 'automatic']:
bpy.ops.object.mode_set(mode='OBJECT')
bpy.ops.object.select_all(action="DESELECT")
self.select()
for obj in self.get_all_objs():
obj.select()
self.armature.select_set(True)
bpy.context.view_layer.objects.active = self.armature
bpy.ops.object.mode_set(mode='POSE')
if weight_distribution == 'envelope':
bpy.ops.object.parent_set(type='ARMATURE_ENVELOPE')
else:
bpy.ops.object.parent_set(type='ARMATURE_AUTO')
bpy.ops.object.mode_set(mode='OBJECT')
elif weight_distribution == 'rigid':
for obj in self.get_all_objs() + [self]:
bpy.ops.object.select_all(action='DESELECT')
obj.blender_obj.parent = self.armature
bpy.context.view_layer.objects.active = obj.blender_obj
mod = obj.blender_obj.modifiers.new("Armature", "ARMATURE")
mod.object = self.armature
obj.blender_obj.vertex_groups.new(name=self.bone.name)
vertices = [v.index for v in obj.blender_obj.data.vertices]
obj.blender_obj.vertex_groups[0].add(vertices, 1.0, 'REPLACE')
bpy.ops.object.select_all(action='DESELECT')
[docs] def _create_ik_bone_controller(self, relative_location: Union[List[float], Vector] = [0., 0., 0.],
use_rotation=True,
chain_length=0) -> (bpy.types.PoseBone, bpy.types.PoseBone, Matrix):
""" Creates an ik bone controller and a corresponding constraint bone for the respective link.
:param relative_location: Relative location of the ik bone controller w.r.t. the bone's location. This can be
used to shift the point of control further away from the end effector.
:param use_rotation: Whether to rotate the child links as well. Defaults to True.
:param chain_length: The number of parent links which are influenced by this ik bone. Defaults to 0 for all
parents.
:return: Constraint and control bone.
"""
bpy.ops.object.select_all(action='DESELECT')
bpy.context.view_layer.objects.active = self.armature
bpy.ops.object.mode_set(mode='EDIT')
self.armature.select_set(True)
edit_bones = self.armature.data.edit_bones
# we need two bones: a controll bone and a constraint bone
# the controll bone will be placed exactly where the current ik bone is
ik_bone_controller = edit_bones.new(self.ik_bone.name + '.controller')
ik_bone_controller.head = self.ik_bone.head + Vector(relative_location)
ik_bone_controller.tail = self.ik_bone.tail + Vector(relative_location)
ik_bone_controller.parent = edit_bones[self.bone.name].parent_recursive[-1]
# the constraint bone will be placed at the head of the ik bone
ik_bone_constraint = edit_bones.new(self.ik_bone.name + '.constraint')
ik_bone_constraint.tail = self.ik_bone.head + Vector(relative_location)
ik_bone_constraint.head = ik_bone_constraint.tail - (self.ik_bone.tail - self.ik_bone.head)
# we need to re-parent the pre- and successor of the constraint bone
ik_bone_constraint.parent = edit_bones[self.ik_bone.name].parent_recursive[0]
edit_bones[self.ik_bone.name].parent = edit_bones[ik_bone_constraint.name]
# add the bones to the link
bpy.ops.object.mode_set(mode='POSE')
self.set_ik_bone_constraint(self.armature.pose.bones.get(self.ik_bone.name + '.constraint'))
self.set_ik_bone_controller(self.armature.pose.bones.get(self.ik_bone.name + '.controller'))
# add ik constraint
set_ik_constraint(self.ik_bone_constraint, self.armature, self.ik_bone_controller.name,
use_rotation=use_rotation, chain_length=chain_length)
if self.get_joint_type() == "revolute":
set_ik_limits_from_rotation_constraint(self.ik_bone_constraint,
constraint=self.bone.constraints["Limit Rotation"])
bpy.ops.object.mode_set(mode='OBJECT')
# determine offset
bpy.context.view_layer.update()
offset = self.ik_bone_constraint.matrix.inverted() @ self.ik_bone_controller.matrix
return self.ik_bone_constraint, self.ik_bone_controller, offset
[docs] def get_fk_ik_mode(self) -> str:
""" Returns the currently selected mode.
:return: One of ["fk", "ik"] denoting forward or inverse kinematics mode.
"""
return self.fk_ik_mode
[docs] def switch_fk_ik_mode(self, mode="fk", keep_pose=True):
""" Switches between forward and inverse kinematics mode. Will do this automatically when switching between e.g.
`set_rotation_euler_fk()` and `set_rotation_euler_ik()`.
:param mode: One of ["fk", "ik"] denoting forward or inverse kinematics mode.
:param keep_pose: If specified, will keep the pose when switching modes. Otherwise will return to the old pose
of the previously selected mode.
"""
if self.bone is None:
return
assert mode in ["fk", "ik"]
if mode == "fk": # turn off copy rotation constraints of fk bone and base bone
if self.get_fk_ik_mode() == "fk":
return True
bpy.context.view_layer.update()
if keep_pose:
self.fk_bone.matrix = self.ik_bone.matrix
if self.joint_type == "revolute":
self.bone.constraints["copy_rotation.fk"].influence = 1.
self.bone.constraints["copy_rotation.ik"].influence = 0.
if "copy_rotation.ik" in self.bone.constraints.keys():
self.bone.constraints["copy_rotation.ik"].influence = 0. # otherwise skip
if self.ik_bone_controller is not None: # switch off ik constraint
self.ik_bone_constraint.constraints["IK"].influence = 0.
self._set_fk_ik_mode(mode="fk")
else: # turn off copy rotation constraints of ik bone and base bone
if self.get_fk_ik_mode() == "ik":
return True
bpy.context.view_layer.update()
if keep_pose:
self.ik_bone.matrix = self.fk_bone.matrix
if self.joint_type == "revolute":
self.bone.constraints["copy_rotation.fk"].influence = 0.
self.bone.constraints["copy_rotation.ik"].influence = 1.
self._set_fk_ik_mode(mode="ik")
if self.ik_bone_controller is not None:
bpy.context.view_layer.update()
location = np.array(self.armature.pose.bones[self.bone.name].head) # or matrix?
print("desired ik location", location)
# location = Vector(np.array(mat)[:3, -1])
self.set_location_ik(location)
if self.ik_bone_constraint is not None:
self.ik_bone_constraint.constraints["IK"].influence = 1.
fk_bone_mat = np.array(self.fk_bone.matrix)
fk_bone_mat[:3, -1] = np.array(self.ik_bone_controller.matrix)[:3, -1]
self.ik_bone_controller.matrix = Matrix(fk_bone_mat)
[docs] def get_joint_rotation(self, frame: int = None) -> float:
""" Get current joint rotation based on euler angles.
:param frame: The desired frame.
:return: Current joint rotation in radians.
"""
if self.bone is None:
return 0.0
pose_bone = self.armature.pose.bones[self.bone.name]
data_bone = self.armature.data.bones[self.bone.name]
with KeyFrame(frame):
M_pose = pose_bone.matrix
M_data = data_bone.matrix_local
# grab the parent's world pose and rest matrices
if data_bone.parent:
M_parent_data = data_bone.parent.matrix_local.copy()
M_parent_pose = pose_bone.parent.matrix.copy()
else:
M_parent_data = Matrix()
M_parent_pose = Matrix()
M1 = M_data.copy()
M1.invert()
M2 = M_parent_pose.copy()
M2.invert()
visual_matrix = M1 @ M_parent_data @ M2 @ M_pose
return visual_matrix.to_quaternion().angle