anime_ikfk_seamless_switch.py

# coding: utf-8
# author DD
# © mihoyo anime
import logging
import sys
import maya.OpenMayaUI as mui
import maya.cmds as cmds
import maya.mel as mel
import pymel.core as pm
import math


try:
    from PySide2 import QtWidgets, QtCore
    from shiboken2 import wrapInstance
    import pyside2uic as uic
except ImportError:
    import PySide.QtGui as QtWidgets
    from PySide import QtCore
    from shiboken import wrapInstance
    import pysideuic as uic

# global params
# fk_shoulder = 'FKScapula'
hand_dict = dict(
    fk_scap='CTL_%sShoulderFK',
    fk_up='CTL_%sArmFK',
    fk_mid='CTL_%sForeArmFK',
    fk_down='CTL_%sHandFK',
    ik_ctrl='CTL_%sArmIK',
    ik_pole='CTL_%sArmPV',
    ik_jnt_up='CTL_%sArmFK',
    ik_jnt_mid='CTL_%sForeArmFK',
    ik_jnt_down='CTL_%sHandFK',
    fk_jnt_up='CTL_%sArmFK',
    fk_jnt_mid='CTL_%sForeArmFK',
    fk_jnt_down='CTL_%sHandFK',
    sw_ctl='CTL_%sArmIK',
    skin_up='CTL_%sArmFK',
    skin_mid='CTL_%sForeArmFK',
    skin_down='CTL_%sHandFK',
)

hand_sw_attr = 'ArmIKBlend'

foot_dict = dict(
    fk_up='CTL_%sUpLegFK',
    fk_mid='CTL_%sLegFK',
    fk_down='CTL_%sFootFK',
    ik_ctrl='CTL_%sFootIK',
    ik_pole='CTL_%sFootPV',
    ik_jnt_up='CTL_%sUpLegFK',
    ik_jnt_mid='CTL_%sLegFK',
    ik_jnt_down='CTL_%sFootFK',
    fk_jnt_up='CTL_%sUpLegFK',
    fk_jnt_mid='CTL_%sLegFK',
    fk_jnt_down='CTL_%sFootFK',
    sw_ctl='CTL_%sFootIK',
    skin_up='CTL_%sUpLegFK',
    skin_mid='CTL_%sLegFK',
    skin_down='CTL_%sFootFK',
)

foot_sw_attr = 'FootIKBlend'
# follow_attr = 'follow'

all_controls_top_group = 'CTL_TransNull'


def get_maya_window():
    main_window_ptr = mui.MQtUtil.mainWindow()
    return wrapInstance(long(main_window_ptr), QtWidgets.QWidget)


def undoable(function):
    # A decorator that will make commands undoable in maya
    def decoratorCode(*args, **kwargs):
        cmds.undoInfo(openChunk=True)
        functionReturn = None
        try:
            functionReturn = function(*args, **kwargs)
        except:
            print(sys.exc_info()[1])
        finally:
            cmds.undoInfo(closeChunk=True)
            return functionReturn
    return decoratorCode


def dotproduct(v1, v2):
    return sum((a*b) for a, b in zip(v1, v2))


def length(v):
    return math.sqrt(dotproduct(v, v))


def distance(p1, p2):
    return math.sqrt((p1[0]-p2[0])**2+(p1[1]-p2[1])**2+(p1[2]-p2[2])**2)


def angle(v1, v2):
    try:
        d = math.degrees(math.acos(dotproduct(v1, v2) / (length(v1) * length(v2))))
        return d
    except:
        return False


class mainWin(QtWidgets.QDialog):
    def __init__(self, parent=get_maya_window()):
        super(mainWin, self).__init__(parent)
        self.setupUi(self)

    def setupUi(self, Dialog):
        Dialog.setObjectName("Dialog")
        # Dialog.setFixedSize(200, 90)
        self.gridLayout = QtWidgets.QGridLayout(Dialog)
        self.gridLayout.setObjectName("gridLayout")
        self.verticalLayout = QtWidgets.QVBoxLayout()
        self.verticalLayout.setObjectName("verticalLayout")
        self.horizontalLayout = QtWidgets.QVBoxLayout()
        self.horizontalLayout.setObjectName("horizontalLayout")
        spacerItem = QtWidgets.QSpacerItem(40, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum)
        self.horizontalLayout.addItem(spacerItem)

        # self.label = QtWidgets.QLabel(Dialog)
        # self.label.setText(u'Author: Dango\n© mihoyo anime')
        self.pushButton = QtWidgets.QPushButton(Dialog)
        self.pushButton.setMinimumSize(QtCore.QSize(100, 50))
        self.pushButton.setObjectName("pushButton")
        self.check_box = QtWidgets.QCheckBox()
        self.check_box.setText(u'AutoKey')
        self.check_box.setChecked(True)
        # self.horizontalLayout.addWidget(self.label)
        self.horizontalLayout.addWidget(self.pushButton)
        self.horizontalLayout.addWidget(self.check_box)
        spacerItem1 = QtWidgets.QSpacerItem(40, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum)
        self.horizontalLayout.addItem(spacerItem1)
        self.verticalLayout.addLayout(self.horizontalLayout)
        self.gridLayout.addLayout(self.verticalLayout, 0, 0, 1, 1)
        QtCore.QMetaObject.connectSlotsByName(Dialog)
        Dialog.setWindowTitle(u"IKFK seamless switch tool by dango © mihoyo anime")
        self.pushButton.setText(u"Switch!")
        self.pushButton.clicked.connect(self.sw)
        self.selected = None
        self.control_value = {}
        self.part = None

    def get_selected(self):
        sel = cmds.ls(sl=1)
        if not sel:
            logging.error(u'Please select any control on arm/foot first!!')
            raise
        self.selected = sel[0]
        self.sw_attr = hand_sw_attr if 'hand' in self.selected.lower() or 'arm' in self.selected.lower() else foot_sw_attr
        return sel[0]

    @property
    def side(self):
        s = 'Left' if 'Left' in self.selected else 'Right'
        return s

    @property
    def namespace(self):
        if ':' in self.selected:
            return ':'.join(self.selected.split(':')[:-1]) + ':'
        return ''

    @property
    def sw_setting(self):
        selected_part = self.selected.replace(self.namespace, '').replace(self.side, '')
        the_set = hand_dict if selected_part in str(hand_dict.values()).replace('%s', '') else foot_dict
        self.part = 'hand'
        if the_set != hand_dict:
            self.part = 'foot'
        return the_set

    def full_name(self, name):
        return self.namespace + self.sw_setting.get(name) % self.side

    def q_r(self, name):
        return cmds.xform(self.full_name(name), q=1, ws=1, ro=1)

    def p_r(self, name, value):
        return cmds.xform(self.full_name(name), ws=1, ro=value)

    def q_t(self, name):
        return cmds.xform(self.full_name(name), q=1, ws=1, t=1)

    def p_t(self, name, value):
        return cmds.xform(self.full_name(name), ws=1, t=value)

    def stick_t(self, orig, dest):
        t = cmds.xform(orig, q=1, ws=1, t=1)
        cmds.xform(dest, ws=1, t=t)

    def stick_r(self, orig, dest):
        t = cmds.xform(orig, q=1, ws=1, ro=1)
        cmds.xform(dest, ws=1, ro=t)

    def get_all_ctrls(self):
        if not cmds.objExists(self.namespace+all_controls_top_group):
            logging.error(u'No %s found!!!' % all_controls_top_group)
            return
        all_nurbs = cmds.listRelatives(self.namespace+all_controls_top_group, c=1, ad=1, type='nurbsCurve')
        all_nurbs_t = [cmds.listRelatives(n, p=1)[0]
                       for n in all_nurbs if pm.PyNode(n).isVisible()]
        return all_nurbs_t

    def deal_with_pole(self):
        if not (cmds.getAttr(self.full_name('fk_mid')+'.r')[0][0] > 1 or cmds.getAttr(self.full_name('fk_mid')+'.r')[0][1] > 1):
            return
        wrist_r = self.q_r('fk_down')
        r = cmds.getAttr(self.full_name('fk_mid')+'.r')
        #  伸直状态
        cmds.setAttr(self.full_name('fk_mid')+'.r', 0, 0, 180)
        p0 = self.q_t('fk_mid')
        p1 = self.q_t('fk_up')
        p2 = self.q_t('fk_down')
        v1 = [p1[0]-p0[0], p1[1]-p0[1], p1[2]-p0[2]]
        v2 = [p2[0]-p0[0], p2[1]-p0[1], p2[2]-p0[2]]
        the_temp_angle = angle(v1, v2) or 0
        # 弯曲
        cmds.setAttr(self.full_name('fk_mid') + '.r', r[0][0], r[0][1], r[0][2])
        p0 = self.q_t('fk_mid')
        p1 = self.q_t('fk_up')
        p2 = self.q_t('fk_down')
        v1 = [p1[0]-p0[0], p1[1]-p0[1], p1[2]-p0[2]]
        v2 = [p2[0]-p0[0], p2[1]-p0[1], p2[2]-p0[2]]
        the_angle = angle(v1, v2)

        cmds.setAttr(self.full_name('fk_mid')+'.r', 0, 0, 180-the_angle-the_temp_angle)

        # 开始转
        new_p2 = self.q_t('fk_down')
        new_v2 = [new_p2[0]-p0[0], new_p2[1]-p0[1], new_p2[2]-p0[2]]
        up_local_rotate_angle = abs(angle(v2, new_v2))
        rv = up_local_rotate_angle #if self.check_box.isChecked() else -up_local_rotate_angle
        fk_up_rx = cmds.getAttr(self.full_name('fk_up')+'.r')[0]
        cmds.setAttr(self.full_name('fk_up')+'.r', fk_up_rx[0]+rv, fk_up_rx[1], fk_up_rx[2])

        new_p3 = self.q_t('fk_down')
        new_v3 = [new_p3[0]-p0[0], new_p3[1]-p0[1], new_p3[2]-p0[2]]
        up_local_rotate_angle_2 = abs(angle(v2, new_v3))

        if 0 < rv < 120:
            if rv < up_local_rotate_angle_2:
                print(u'1.转过了！')
                cmds.setAttr(self.full_name('fk_up')+'.r', fk_up_rx[0]-2*rv, fk_up_rx[1], fk_up_rx[2])
        if abs(up_local_rotate_angle_2 - (360 - (2*rv))) < 1:
            print(u'2.转过了！')
            cmds.setAttr(self.full_name('fk_up')+'.r', fk_up_rx[0]-2*rv, fk_up_rx[1], fk_up_rx[2])
        # cmds.rotate(rv, 0, 0, self.full_name('fk_up'), objectSpace=1, relative=1)

        # new_p0 = self.q_t('fk_mid')
        # new_p2 = self.q_t('fk_down')
        # new_v2 = [new_p2[0]-new_p0[0], new_p2[1]-new_p0[1], new_p2[2]-new_p0[2]]
        # new_up_local_rotate_angle = angle(v2, new_v2)
        # reverse = False
        # if 0 < up_local_rotate_angle < 90:
        #     if new_up_local_rotate_angle > up_local_rotate_angle:
        #         reverse = True
        # elif 90 < up_local_rotate_angle < 180:
        # print distance(p2, new_p2)
        # if distance(p2, new_p2) > 0.5:
        # if self.check_box.isChecked():
        # if abs(new_up_local_rotate_angle - up_local_rotate_angle) >= 5:
                # or abs(360.0-new_up_local_rotate_angle - 2*up_local_rotate_angle) <= 5:
            # cmds.rotate(-2*up_local_rotate_angle, 0, 0, self.full_name('fk_up'), objectSpace=1, relative=1)

        self.p_r('fk_down', wrist_r)
        print('elbow fixed!')

    def reset_ctrls(self):
        self.control_value = {}
        att_list = ['.rotateX', '.rotateY', '.rotateZ', '.translateX', '.translateY', '.translateZ']
        all_ctrls = self.get_all_ctrls()
        for ctl in all_ctrls:
            for at in att_list:
                self.control_value[ctl+at] = cmds.getAttr(ctl+at)
        for full_attr in self.control_value.keys():
            try:
                cmds.setAttr(full_attr, 0)
            except:
                pass
        pass

    def restore_contrls(self):
        for a, v in self.control_value.items():
            try:
                cmds.setAttr(a, v)
            except:
                pass
        self.control_value = {}

    def ik_2_fk(self):
        for ss in ['up', 'mid', 'down']:
            self.p_r('fk_%s' % ss, self.q_r('ik_jnt_%s' % ss))
        cmds.setAttr(self.full_name('sw_ctl') + '.' + self.sw_attr, 0)
        pass

    def fk_2_ik(self):
        self.deal_with_pole()
        self.reset_ctrls()
        cmds.setAttr(self.full_name('sw_ctl') + '.' + self.sw_attr, 1)
        cmds.select(self.full_name('ik_ctrl'), r=1)
        p_temp = cmds.duplicate(rr=1)
        cmds.parentConstraint(self.full_name('fk_down'), p_temp, mo=1)
        locator = cmds.spaceLocator()
        self.stick_t(self.full_name('ik_pole'), locator)
        cmds.parentConstraint(self.full_name('fk_up'), locator, mo=1)
        self.restore_contrls()
        cmds.setAttr(self.full_name('sw_ctl') + '.' + self.sw_attr, 0)
        self.stick_t(locator, self.full_name('ik_pole'))
        # if not self.part == 'hand':
        self.stick_t(p_temp, self.full_name('ik_ctrl'))
        self.stick_r(p_temp, self.full_name('ik_ctrl'))
        cmds.delete(p_temp)
        cmds.delete(locator)
        cmds.setAttr(self.full_name('sw_ctl') + '.' + self.sw_attr, 1)
        pass

    @undoable
    def sw(self):
        self.get_selected()
        if self.check_box.isChecked():
            ct = cmds.currentTime(q=1)
            cmds.currentTime(ct-1)
            for each in ['sw_ctl', 'ik_ctrl', 'ik_pole', 'fk_up', 'fk_mid', 'fk_down']:
                cmds.select(self.full_name(each), r=1)
                mel.eval('SetKey;')
            cmds.currentTime(ct) 

        if cmds.getAttr(self.full_name('sw_ctl') + '.' + self.sw_attr) > 0:
            self.ik_2_fk()
        else:
            self.fk_2_ik()

        logging.info(u'Success！')
        if self.check_box.isChecked():
            for each in ['sw_ctl', 'ik_ctrl', 'ik_pole', 'fk_up', 'fk_mid', 'fk_down']:
                cmds.select(self.full_name(each), r=1)
                mel.eval('SetKey;')
        cmds.select(self.selected)


if __name__ == '__main__':
    win = mainWin()
    win.show()