Autodesk Maya tool with user interface to fit a rigid transform to a set of locators.

fit_rigid_3d_transform.py

import numpy as np
from math import sqrt
import maya.cmds as cmds
from maya.api import OpenMaya


def str_to_dagnode(name):
    '''
    Returns a Maya dagNode from an input string.
    :param name: name of the object in Maya
    '''
    selection_list = OpenMaya.MSelectionList()
    try:
        selection_list.add(name)
    except RuntimeError:
        return None
    dagNode = OpenMaya.MDagPath()
    dagNode = selection_list.getDagPath(0)
    return dagNode

def mat3_to_mat4(mat):
    '''Converts 3x3 matrix into 4x4'''
    mat3 = mat
    if isinstance(mat, list):
        mat3 = np.array(mat)
    mat4 = np.zeros((4, 4))
    mat4[:3, :3] = mat3
    mat4[3][3] = 1.
    if isinstance(mat, list):
        return mat4.tolist()
    return mat4

def isRotationMatrix(R):
    '''Checks if a matrix is a valid rotation matrix.'''
    Rt = np.transpose(R)
    shouldBeIdentity = np.dot(Rt, R)
    I = np.identity(3, dtype = R.dtype)
    n = np.linalg.norm(I - shouldBeIdentity)
    return n < 1e-6

def rotationMatrixToEulerAngles(R):
    '''
    Calculates rotation matrix to euler angles
    The result is the same as MATLAB except the order
    of the euler angles ( x and z are swapped ).
    '''
    assert(isRotationMatrix(R))
     
    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
     
    singular = sy < 1e-6
 
    if  not singular :
        x = math.atan2(R[2,1] , R[2,2])
        y = math.atan2(-R[2,0], sy)
        z = math.atan2(R[1,0], R[0,0])
    else :
        x = math.atan2(-R[1,2], R[1,1])
        y = math.atan2(-R[2,0], sy)
        z = 0
 
    return [z, y, x]

def center_on_points(mesh, points):
    '''
    Centers pivot point based on selected mesh components (vertice,
    edge, face)
    '''
    piv = np.mean(points, axis=0)
    pivot = [piv.tolist()[0], piv.tolist()[1], piv.tolist()[2]]
    cmds.xform(mesh, pivots=pivot, worldSpace=True)

def ui_add_selected_pts(sel_pts, col):
    '''Adds selected points to table widget'''
    last_row_num = cmds.scriptTable('table', query=True, rows=True) - 1
    
    # Clean rows before adding new items
    for row in range(last_row_num):
        cmds.scriptTable(
            TABLE, edit=1, cellIndex=(row + 1, col), cellValue='')
    for row in range(last_row_num):
        cmds.scriptTable(
            TABLE, edit=1, cellIndex=(row + 1, col), cellValue='')
    # Create enough rows to allocate selected points
    if last_row_num < len(sel_pts):
        for n in range(len(sel_pts) - 1):
            last_row_num = cmds.scriptTable('table', query=True, rows=True)
            cmds.scriptTable('table', edit=True, insertRow=last_row_num)
    # Add points to table
    for pt in range(len(sel_pts)):
        cmds.scriptTable(
            TABLE, edit=1, cellIndex=(pt+1, col), cellValue=sel_pts[pt])

def ui_add_src_pts(args):
    sel_pts = cmds.ls(sl=1, flatten=1)
    ui_add_selected_pts(sel_pts, 1)

def ui_add_tgt_pts(args):
    sel_pts = cmds.ls(sl=1, flatten=1)
    ui_add_selected_pts(sel_pts, 2)

def ui_get_points():
    points_a = []
    points_b = []
    max_rows = cmds.scriptTable('table', q=1, rows=1)
    for r in range(1, max_rows):
        pta = cmds.scriptTable('table', q=1, cellIndex=(r, 1), cellValue=1)
        if isinstance(pta, list):
            pta = pta[0]
        if pta:
            ppos_a = cmds.pointPosition(pta, w=1)
            points_a.append(ppos_a)
        ptb = cmds.scriptTable('table', q=1, cellIndex=(r, 2), cellValue=1)
        if isinstance(ptb, list):
            ptb = ptb[0]
        if ptb:
            ppos_b = cmds.pointPosition(ptb, w=1)
            points_b.append(ppos_b)
    return(points_a, points_b)

def get_target_mesh():
    '''Figure out the transform mesh from selected points'''
    src_mesh = cmds.scriptTable(
        'table', q=1, cellIndex=(1, 1), cellValue=1)[0]
    transform = None
    if cmds.objectType(src_mesh) == 'mesh':
        parent_shape = cmds.listRelatives(src_mesh, allParents=1)
        if 'Shape' in parent_shape[0]:
            transform = cmds.listRelatives(parent_shape, allParents=1)
        else:
            transform = parent_shape
    elif cmds.objectType(src_mesh) == 'transform':
        transform = src_mesh
    if not transform:
        raise Exception('Unable to find transform mesh from selected source points')
    if isinstance(transform, list):
        transform = transform[0]
    return transform

def ui_edit_cell(row, column, value):
    return 1

def ui_open_dialog():
    '''Maya Dialog Window'''
    global TABLE
    # Remove existing windows first
    for win in cmds.lsUI(windows=1):
        if cmds.window(win, q=1, title=1) == 'Fit Transform':
            cmds.deleteUI(win)

    window_ = cmds.window(widthHeight=(410, 350), title='Fit Transform')
    form = cmds.formLayout()

    TABLE = cmds.scriptTable(
        'table',
        rows=1,
        columns=2,
        columnWidth=([1, 197], [2, 197]),
        label=[(1, "Source points"), (2, "Target points")],
        cellChangedCmd=ui_edit_cell)
    src_button = cmds.button(label="Add Selected Source", command=ui_add_src_pts)
    target_button = cmds.button(label="Add Selected Target", command=ui_add_tgt_pts)
    apply_button = cmds.button(label="Apply", command=fit_pose)
    cmds.formLayout(
        form,
        edit=True,
        attachForm=[(TABLE, 'top', 0), (TABLE, 'left', 0),
                    (TABLE, 'right', 0), (src_button, 'bottom', 25),
                    (src_button, 'left', 0), (target_button, 'bottom', 25),
                    (target_button, 'right', 0), (apply_button, 'bottom', 0),
                    (apply_button, 'left', 0)],
        attachControl=(TABLE, 'bottom', 0, src_button),
        attachNone=[(src_button, 'top'), (target_button, 'top'),
                    (apply_button, 'top')],
        attachPosition=[(src_button, 'right', 0, 50),
                        (target_button, 'left', 0, 50), 
                        (apply_button, 'right', 0, 100)])
    cmds.showWindow(window_)
    return TABLE

def singular_value_decomposition(A, B):
    '''
    Finds optimal rotation and translation between 3D points,
    using Singular Value Decomposition (SVD).
    
    :param A: a Nx3 numpy matrix representing the source 3D points
    :param B: a Nx3 numpy matrix representing the target 3D points

    :returns R: 3x3 rotation matrix
    :returns t: 3x1 column vector representing translation

    See: https://nghiaho.com/?page_id=671
         https://github.com/demotu/BMC/blob/master/functions/svdt.py
    '''
    if len(A) != len(B):
        msg = "Number of point count must match!"
        cmds.confirmDialog(title='Error', message=msg, button='OK')
        raise Exception("Number of points between A and B must match!")

    N = A.shape[0]  # total points

    centroid_A = np.mean(A, axis=0)
    centroid_B = np.mean(B, axis=0)
    
    # center the points at the origin
    AA = A - np.tile(centroid_A, (N, 1))
    BB = B - np.tile(centroid_B, (N, 1))

    # H is the familiar covariance matrix
    H = np.transpose(AA) * BB

    # Singular Value Decomposition (SVD)
    U, S, Vt = np.linalg.svd(H)

    # Rotation Matrix
    R = Vt.T * U.T

    # here we take care of a special reflection case, which
    # yields a nonsense negative value sometimes
    if np.linalg.det(R) < 0:
       Vt[2,:] *= -1
       R = Vt.T * U.T

    # Find the translation vector
    t = centroid_B.T - np.dot(R, centroid_A.T)

    # Find Root-mean-squared error
    A2 = (R*A.T) + np.tile(t, (1, N))
    A2 = A2.T
    err = A2 - B
    err = np.multiply(err, err)
    err = np.sum(err)
    rmse = np.sqrt(err/N)

    return (R, t, rmse)

def fit_pose(args):
    '''
    Applies the rotation matrix(R) and translation vector(t)
    to the target Maya transform mesh.
    '''

    target_mesh = get_target_mesh()
    points_a, points_b = ui_get_points()
    if not points_a or not points_b:
        print('No points defined')

    cmds.makeIdentity(target_mesh, apply=True, t=1, r=1, s=1, n=2 )
    R, t, rmse = singular_value_decomposition(np.mat(points_a), np.mat(points_b))
    print('RMSE: %s' % rmse)

    # Finally, fit target mesh to source points
    dag_node = str_to_dagnode(target_mesh)
    mfnxform = OpenMaya.MFnTransform(dag_node)
    rotation_maya_matrix = OpenMaya.MMatrix(mat3_to_mat4(R.transpose()))
    rotation_xform = OpenMaya.MTransformationMatrix(rotation_maya_matrix)
    import pdb; pdb.set_trace()
    mfnxform.setTransformation(rotation_xform)

    # Set pivot point before applying translation
    center_on_points(target_mesh, points_a)
    translation = [x[0] for x in t.tolist()]
    cmds.xform(target_mesh, t=translation, ws=1)


ui_open_dialog()