MetaBox/Scripts/Animation/epic_pose_wrangler/v2/model/utils.py
2025-01-14 03:06:35 +08:00

312 lines
11 KiB
Python

# Copyright Epic Games, Inc. All Rights Reserved.
import math
import traceback
from maya import OpenMaya, cmds
from epic_pose_wrangler.log import LOG
from epic_pose_wrangler.v2.model import exceptions
# NOTE: MTransformationMatrix & MEulerRotation have different values for the same axis.
XFORM_ROTATION_ORDER = {
'xyz': OpenMaya.MTransformationMatrix.kXYZ,
'yzx': OpenMaya.MTransformationMatrix.kYZX,
'zxy': OpenMaya.MTransformationMatrix.kZXY,
'xzy': OpenMaya.MTransformationMatrix.kXZY,
'yxz': OpenMaya.MTransformationMatrix.kYXZ,
'zyx': OpenMaya.MTransformationMatrix.kZYX
}
EULER_ROTATION_ORDER = {
'xyz': OpenMaya.MEulerRotation.kXYZ,
'yzx': OpenMaya.MEulerRotation.kYZX,
'zxy': OpenMaya.MEulerRotation.kZXY,
'xzy': OpenMaya.MEulerRotation.kXZY,
'yxz': OpenMaya.MEulerRotation.kYXZ,
'zyx': OpenMaya.MEulerRotation.kZYX
}
def compose_matrix(position, rotation, rotation_order='xyz'):
"""
Compose a 4x4 matrix with given transformation.
>>> compose_matrix((0.0, 0.0, 0.0), (90.0, 0.0, 0.0))
[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
"""
# create rotation ptr
rot_script_util = OpenMaya.MScriptUtil()
rot_script_util.createFromDouble(*[deg * math.pi / 180.0 for deg in rotation])
rot_double_ptr = rot_script_util.asDoublePtr()
# construct transformation matrix
xform_matrix = OpenMaya.MTransformationMatrix()
xform_matrix.setTranslation(OpenMaya.MVector(*position), OpenMaya.MSpace.kTransform)
xform_matrix.setRotation(rot_double_ptr, XFORM_ROTATION_ORDER[rotation_order], OpenMaya.MSpace.kTransform)
matrix = xform_matrix.asMatrix()
return [matrix(m, n) for m in range(4) for n in range(4)]
def decompose_matrix(matrix, rotation_order='xyz'):
"""
Decomposes a 4x4 matrix into translation and rotation.
>>> decompose_matrix([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
((0.0, 0.0, 0.0), (90.0, 0.0, 0.0))
"""
mmatrix = OpenMaya.MMatrix()
OpenMaya.MScriptUtil.createMatrixFromList(matrix, mmatrix)
# create transformation matrix
xform_matrix = OpenMaya.MTransformationMatrix(mmatrix)
# get translation
translation = xform_matrix.getTranslation(OpenMaya.MSpace.kTransform)
# get rotation
# @ref: https://github.com/LumaPictures/pymel/blob/master/pymel/core/datatypes.py
# The apicls getRotation needs a "RotationOrder &" object, which is impossible to make in python...
euler_rotation = xform_matrix.eulerRotation()
euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
rotation = euler_rotation.asVector()
return (
(translation.x, translation.y, translation.z),
(rotation.x * 180.0 / math.pi, rotation.y * 180.0 / math.pi, rotation.z * 180.0 / math.pi)
)
def euler_to_quaternion(rotation, rotation_order='xyz'):
"""
Returns Euler Rotation as Quaternion
>>> euler_to_quaternion((90, 0, 0))
(0.7071, 0.0, 0.0, 0.70710))
"""
euler_rotation = OpenMaya.MEulerRotation(
rotation[0] * math.pi / 180.0,
rotation[1] * math.pi / 180.0,
rotation[2] * math.pi / 180.0
)
euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
quat = euler_rotation.asQuaternion()
return quat.x, quat.y, quat.z, quat.w
def quaternion_to_euler(rotation, rotation_order='xyz'):
"""
Returns Quaternion Rotation as Euler
quaternion_to_euler((0.7071, 0.0, 0.0, 0.70710))
(90, 0, 0)
"""
quat = OpenMaya.MQuaternion(*rotation)
euler_rotation = quat.asEulerRotation()
euler_rotation.reorderIt(EULER_ROTATION_ORDER[rotation_order])
return euler_rotation.x * 180.0 / math.pi, euler_rotation.y * 180.0 / math.pi, euler_rotation.z * 180.0 / math.pi
def is_connected_to_array(attribute, array_attr):
"""
Check if the attribute is connected to the specified array
:param attribute :type str: attribute
:param array_attr :type str array attribute
:return :type int or None: int of index in the array or None
"""
try:
indices = cmds.getAttr(array_attr, multiIndices=True) or []
except ValueError:
return None
for i in indices:
attr = '{from_attr}[{index}]'.format(from_attr=array_attr, index=i)
if attribute in (cmds.listConnections(attr, plugs=True) or []):
return i
return None
def get_next_available_index_in_array(attribute):
# Get the next available index
indices = cmds.getAttr(attribute, multiIndices=True) or []
i = 0
for index in indices:
if index != i and i not in indices:
indices.append(i)
i += 1
indices.sort()
attrs = ['{from_attr}[{index}]'.format(from_attr=attribute, index=i) for i in indices]
connections = [cmds.listConnections(attr, plugs=True) or [] for attr in attrs]
target_index = len(indices)
for index, conn in enumerate(connections):
if not conn:
target_index = index
break
return target_index
def message_connect(from_attribute, to_attribute, in_array=False, out_array=False):
"""
Create and connect a message attribute between two nodes
"""
# Generate the object and attr names
from_object, from_attribute_name = from_attribute.split('.', 1)
to_object, to_attribute_name = to_attribute.split('.', 1)
# If the attributes don't exist, create them
if not cmds.attributeQuery(from_attribute_name, node=from_object, exists=True):
cmds.addAttr(from_object, longName=from_attribute_name, attributeType='message', multi=in_array)
if not cmds.attributeQuery(to_attribute_name, node=to_object, exists=True):
cmds.addAttr(to_object, longName=to_attribute_name, attributeType='message', multi=out_array)
# Check that both attributes, if existing are message attributes
for a in (from_attribute, to_attribute):
if cmds.getAttr(a, type=1) != 'message':
raise exceptions.MessageConnectionError(
'Message Connect: Attribute {attr} is not a message attribute. CONNECTION ABORTED.'.format(
attr=a
)
)
# Connect up the attributes
try:
if in_array:
from_attribute = "{from_attribute}[{index}]".format(
from_attribute=from_attribute,
index=get_next_available_index_in_array(from_attribute)
)
if out_array:
to_attribute = "{to_attribute}[{index}]".format(
to_attribute=to_attribute,
index=get_next_available_index_in_array(to_attribute)
)
return cmds.connectAttr(from_attribute, to_attribute, force=True)
except Exception as e:
LOG.error(traceback.format_exc())
return False
def connect_attr(from_attr, to_attr):
if not cmds.isConnected(from_attr, to_attr):
cmds.connectAttr(from_attr, to_attr)
def get_attr(attr_name, as_value=True):
"""
Get the specified attribute
:param attr_name :type str: attribute name i.e node.translate
:param as_value :type bool: return as value or connected plug name
:return :type list or any: either returns a list of connections or the value of the attribute
"""
# Check if the attribute is connected
connections = cmds.listConnections(attr_name, plugs=True)
if connections and not as_value:
# If the attribute is connected and we don't want the value, return the connections
return connections
elif as_value:
# Return the value
return cmds.getAttr(attr_name)
def get_attr_array(attr_name, as_value=True):
"""
Get the specified array attr
:param attr_name :type str: attribute name i.e node.translate
:param as_value :type bool: return as value or connected plug name
:return :type list or any: either returns a list of connections or the value of the attribute
"""
# Get the number of indices in the array
indices = cmds.getAttr(attr_name, multiIndices=True) or []
# Empty list to store the connected plugs
connected_plugs = []
# Empty list to store values
values = []
# Iterate through the indices
for i in indices:
# Get all the connected plugs for this index
connections = cmds.listConnections('{attr_name}[{index}]'.format(attr_name=attr_name, index=i), plugs=True)
# If we want the plugs and not values, store connections
if connections and not as_value:
connected_plugs.extend(connections)
# If we want values, get the value at the index
elif as_value:
values.append(cmds.getAttr('{attr_name}[{index}]'.format(attr_name=attr_name, index=i)))
# Return plugs or values, depending on which one has data
return connected_plugs or values
def set_attr_or_connect(source_attr_name, value=None, attr_type=None, output=False):
"""
Set an attribute or connect it to another attribute
:param source_attr_name :type str: attribute name
:param value : type any: value to set the attribute to
:param attr_type :type str: name of the attribute type i.e matrix
:param output :type bool: is this plug an output (True) or input (False)
"""
# Type conversion from maya: python
attr_types = {
'matrix': list
}
# Check if we have a matching type
matching_type = attr_types.get(attr_type, None)
# If we have a matching type and the value matches that type, set the attr
if matching_type is not None and isinstance(value, matching_type):
cmds.setAttr(source_attr_name, value, type=attr_type)
# If the value is a string and no type is matched, we want to connect the attributes
elif isinstance(value, str):
try:
# Connect from left->right depending on if the source is output or input
if output:
if not cmds.isConnected(source_attr_name, value):
cmds.connectAttr(source_attr_name, value)
else:
if not cmds.isConnected(value, source_attr_name):
cmds.connectAttr(value, source_attr_name)
except Exception as e:
raise exceptions.PoseWranglerAttributeError(
"Unable to {direction} {input} to '{output}'".format(
direction="connect" if value else "disconnect",
input=source_attr_name if output else value,
output=value if output else source_attr_name
)
)
else:
cmds.setAttr(source_attr_name, value)
def disconnect_attr(attr_name, array=False):
"""
Disconnect the specified attribute
:param attr_name :type str: attribute name to disconnect
:param array :type bool: is this attribute an array?
"""
attrs = []
# If we are disconnecting an array, get the names of all the attributes
if array:
attrs.extend(cmds.getAttr(attr_name, multiIndices=True) or [])
# Otherwise append the attr name specified
else:
attrs.append(attr_name)
# Iterate through all the attrs listed
for attr in attrs:
# Find their connections and disconnect them
for plug in cmds.listConnections(attr, plugs=True) or []:
cmds.disconnectAttr(attr, plug)
def get_selection(_type=""):
"""
Returns the current selection
"""
return cmds.ls(selection=True, type=_type)
def set_selection(selection_list):
"""
Sets the active selection
"""
cmds.select(selection_list, replace=True)