MetaBox/Scripts/Animation/epic_pose_wrangler/v1/poseWranglerUI.py
2025-01-14 03:05:57 +08:00

659 lines
22 KiB
Python

# Copyright Epic Games, Inc. All Rights Reserved.
from PySide2 import QtWidgets
from PySide2 import QtCore
from PySide2 import QtUiTools
import os
import maya.cmds as cmds
import maya.OpenMaya as OpenMaya
from maya.app.general.mayaMixin import MayaQWidgetDockableMixin
from epic_pose_wrangler.log import LOG
from epic_pose_wrangler.view import log_widget
from epic_pose_wrangler.v1 import poseWrangler, upgrade
from epic_pose_wrangler.v1 import palette
class EventUpgrade(QtCore.QObject):
upgrade = QtCore.Signal(str)
class PoseWrangler(MayaQWidgetDockableMixin, QtWidgets.QMainWindow):
"""
class for the pose wranglerUI
"""
ui_name = "poseWranglerWindow"
def __init__(self, parent=None):
super(PoseWrangler, self).__init__(parent)
self.event_upgrade_dispatch = EventUpgrade()
# Load the dependent plugin
plugin_versions = [{
"name": "MayaUE4RBFPlugin_{}".format(cmds.about(version=True)),
"node_name": "UE4RBFSolverNode"
},
{
"name": "MayaUE4RBFPlugin{}".format(cmds.about(version=True)),
"node_name": "UE4RBFSolverNode"
},
{"name": "MayaUERBFPlugin".format(cmds.about(version=True)), "node_name": "UERBFSolverNode"}]
QtCore.QSettings.setPath(QtCore.QSettings.IniFormat, QtCore.QSettings.UserScope, os.environ['LOCALAPPDATA'])
self._settings = QtCore.QSettings(
QtCore.QSettings.IniFormat, QtCore.QSettings.UserScope, "Epic Games",
"PoseWrangler"
)
self._settings.setFallbacksEnabled(False)
for plugin_version in plugin_versions:
plugin_name = plugin_version['name']
node_name = plugin_version['node_name']
if cmds.pluginInfo(plugin_name, q=True, loaded=True):
self._node_name = node_name
break
else:
try:
cmds.loadPlugin(plugin_name)
self._node_name = node_name
break
except RuntimeError as e:
pass
else:
raise RuntimeError("Unable to load valid RBF plugin version")
# Load the UI file
file_path = os.path.dirname(__file__) + "/poseWranglerUI.ui"
if os.path.exists(file_path):
ui_file = QtCore.QFile(file_path)
# Attempt to open and load the UI
try:
ui_file.open(QtCore.QFile.ReadOnly)
loader = QtUiTools.QUiLoader()
self.win = loader.load(ui_file)
finally:
# Always close the UI file regardless of loader result
ui_file.close()
else:
raise ValueError('UI File does not exist on disk at path: {}'.format(file_path))
self.setWindowTitle("Pose Wrangler")
# Embed the UI window inside this widget
self.setCentralWidget(self.win)
# buttons
self.win.add_pose_BTN.pressed.connect(self.add_pose)
# refresh the driver cmd list
self.load_drivers()
# hook up utils UI
self.win.bake_poses_BTN.pressed.connect(self.bake_poses)
self.win.pose_LIST.itemSelectionChanged.connect(self.pose_changed)
self.win.edit_pose_BTN.pressed.connect(self.edit_pose)
self.win.delete_pose_BTN.pressed.connect(self.delete_pose)
self.win.driver_LIST.itemSelectionChanged.connect(self.driver_changed)
self.win.driver_LIST.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
self.win.driver_LIST.customContextMenuRequested.connect(self.driver_popup)
self.win.driven_transforms_LIST.itemSelectionChanged.connect(self.driven_changed)
self.win.select_driver_BTN.pressed.connect(self.select_driver)
self.win.add_driven_BTN.pressed.connect(self.add_new_driven)
self.win.mirror_pose_BTN.pressed.connect(self.mirror_pose)
self.win.copy_driven_trs_BTN.pressed.connect(self.copy_driven_trs)
self.win.paste_driven_trs_BTN.pressed.connect(self.paste_driven_trs)
self.win.create_driver_BTN.pressed.connect(self.create_driver)
self.win.delete_driver_BTN.pressed.connect(self.delete_driver)
# setup the 'driving enabled' check box functionality
self.win.toggle_edit_BTN.clicked.connect(self.toggle_edit)
self.win.refresh_BTN.clicked.connect(self.load_drivers)
self.win.driver_transform_LINE.setReadOnly(True)
self.win.upgrade_scene_BTN.pressed.connect(self._upgrade_scene)
self._log_widget = log_widget.LogWidget()
self.addDockWidget(QtCore.Qt.BottomDockWidgetArea, self._log_widget.log_dock)
LOG.addHandler(self._log_widget)
# refresh the tree
self.load_poses()
self.set_stylesheet()
self._driver = None
def set_stylesheet(self):
"""set the theming and styling here"""
self.setStyleSheet(palette.getPaletteString())
def driver_popup(self, point):
"""add right click menu"""
selected_solvers = self.get_selected_solvers()
solver = None
if selected_solvers:
solver = selected_solvers[-1]
"""--------------------------EDIT----------------------------"""
edit_action = QtWidgets.QAction("Edit", self)
edit_action.triggered.connect(self.edit_driver)
enable_action = QtWidgets.QAction("Finish Editing (Enable)", self)
enable_action.triggered.connect(self.enable_driver)
"""--------------------------SELECTION----------------------------"""
select_driver_action = QtWidgets.QAction("Select Driver Transform", self)
select_driver_action.triggered.connect(self.select_driver)
select_solver_action = QtWidgets.QAction("Select Solver Node", self)
select_solver_action.triggered.connect(self.select_solver)
select_driven_action = QtWidgets.QAction("Select Driven Transform(s)", self)
select_driven_action.triggered.connect(self.select_driven)
"""--------------------------MIRROR----------------------------"""
mirror_driver_action = QtWidgets.QAction("Mirror Selected Drivers", self)
mirror_driver_action.triggered.connect(self.mirror_driver)
"""----------------------IMPORT/EXPORT----------------------------"""
export_driver_action = QtWidgets.QAction("Export Selected Drivers", self)
export_driver_action.triggered.connect(self.export_driver)
import_driver_action = QtWidgets.QAction("Import Driver(s)", self)
import_driver_action.triggered.connect(self.import_drivers)
export_all_action = QtWidgets.QAction("Export All Drivers", self)
export_all_action.triggered.connect(self.export_all)
"""--------------------------ADD----------------------------"""
add_driven_action = QtWidgets.QAction("Add Driven Transform(s)", self)
add_driven_action.triggered.connect(self.add_new_driven)
"""--------------------------UTILITIES----------------------------"""
zero_base_pose_action = QtWidgets.QAction("Zero Base Pose Transforms", self)
zero_base_pose_action.triggered.connect(self.zero_base_poses)
menu = QtWidgets.QMenu("Options:", self.win.driver_LIST)
select_menu = QtWidgets.QMenu("Select:", menu)
select_menu.addAction(select_driver_action)
select_menu.addAction(select_driven_action)
select_menu.addAction(select_solver_action)
mirror_menu = QtWidgets.QMenu("Mirror:", menu)
mirror_menu.addAction(mirror_driver_action)
import_export_menu = QtWidgets.QMenu("Import/Export:", menu)
if solver:
import_export_menu.addAction(export_driver_action)
import_export_menu.addAction(export_all_action)
import_export_menu.addSeparator()
import_export_menu.addAction(import_driver_action)
add_menu = QtWidgets.QMenu("Add:", menu)
add_menu.addAction(add_driven_action)
utilities_menu = QtWidgets.QMenu("Utilities:", menu)
utilities_menu.addAction(zero_base_pose_action)
if solver:
# if the solver is enabled show the edit and otherwise show the enable
if solver.is_enabled:
menu.addAction(edit_action)
else:
menu.addAction(enable_action)
menu.addSeparator()
menu.addMenu(select_menu)
menu.addMenu(add_menu)
menu.addMenu(mirror_menu)
menu.addMenu(utilities_menu)
menu.addMenu(import_export_menu)
menu.popup(self.win.driver_LIST.mapToGlobal(point))
def get_selected_solvers(self):
"""gets the selected solvers"""
selected_items = self.win.driver_LIST.selectedItems()
selected_solvers = []
if selected_items:
for item in selected_items:
solver = item.data(QtCore.Qt.UserRole)
selected_solvers.append(solver)
return selected_solvers
def get_selected_poses(self):
"""gets the selected poses"""
selected_items = self.win.pose_LIST.selectedItems()
selected_poses = []
if selected_items:
for item in selected_items:
pose = item.text()
selected_poses.append(pose)
return selected_poses
def add_pose(self):
"""add a new pose to the driver"""
selected_solvers = self.get_selected_solvers()
if not selected_solvers:
OpenMaya.MGlobal.displayError('PoseWrangler: You must have a driver selected!')
return
solver = selected_solvers[-1]
pose_name, ok = QtWidgets.QInputDialog.getText(self, 'text', 'Pose Name:')
if pose_name:
solver.add_pose(pose_name)
# self.win.edit_pose_BTN.setEnabled(True)
self.load_poses()
else:
cmds.warning('PoseWrangler: You must enter a pose name to add a pose.')
def edit_pose(self):
"""updates the current pose"""
poses = self.get_selected_poses()
if not poses:
return
pose = poses[-1]
selected_solvers = self.get_selected_solvers()
if not selected_solvers:
OpenMaya.MGlobal.displayError('PoseWrangler: You must have a driver selected!')
return
solver = selected_solvers[-1]
solver.update_pose(pose)
def delete_pose(self):
"""deletes the selected pose"""
poses = self.get_selected_poses()
if not poses:
return
pose = poses[-1]
selected_solvers = self.get_selected_solvers()
if not selected_solvers:
OpenMaya.MGlobal.displayError('PoseWrangler: You must have a driver selected!')
return
solver = selected_solvers[-1]
solver.delete_pose(pose)
self.load_poses()
def driver_changed(self):
"""function that gets called with the driver/solver is clicked"""
selected_solvers = self.get_selected_solvers()
self.win.driven_transforms_LIST.clear()
if selected_solvers:
solver = selected_solvers[-1]
driven_transforms = solver.driven_transforms
if driven_transforms:
for transform in driven_transforms:
item = QtWidgets.QListWidgetItem(transform)
self.win.driven_transforms_LIST.addItem(item)
else:
self.win.driver_transform_LINE.setText("")
self.refresh_ui_state()
self.load_poses()
def driven_changed(self):
"""select the driven transforms when picked in the UI"""
selected_items = self.win.driven_transforms_LIST.selectedItems()
cmds.select(cl=1)
for item in selected_items:
cmds.select(item.text(), add=1)
def select_driver(self):
"""selects the driver(s)"""
cmds.select(cl=1)
selected_solvers = self.get_selected_solvers()
for solver in selected_solvers:
cmds.select(solver.driving_transform, add=1)
def select_driven(self):
"""selects the driven"""
cmds.select(cl=1)
selected_solvers = self.get_selected_solvers()
for solver in selected_solvers:
cmds.select(solver.driven_transforms, add=1)
def select_solver(self):
"""selects the solver DG node"""
cmds.select(cl=1)
selected_solvers = self.get_selected_solvers()
for solver in selected_solvers:
cmds.select(solver.name, add=1)
def bake_poses(self):
"""bakes the poses to the timeline"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.bake_poses_to_timeline()
def add_new_driven(self):
"""adds new driven transforms to the selected drivers"""
selected_items = self.win.driver_LIST.selectedItems()
sl = cmds.ls(sl=1)
if not sl:
OpenMaya.MGlobal.displayError("No driven object selected, please select a driven object to add")
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
for s in sl:
solver.add_driven(s)
self.load_drivers(selected_items[-1].text())
def edit_driver(self):
"""set the drivers into edit mode"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.is_driving(False)
self.refresh_ui_state()
def enable_driver(self):
"""enables the drivers into finishes edit mode"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.is_driving(True)
self.refresh_ui_state()
def toggle_edit(self):
"""toggle the edit"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
solver = selected_solvers[-1]
if solver.is_enabled:
solver.is_driving(False)
else:
solver.is_driving(True)
self.refresh_ui_state()
def refresh_ui_state(self):
"""refresh the UI state"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
solver = selected_solvers[-1]
self.win.driver_transform_LINE.setText(solver.driving_transform)
# enable buttons that should be available when the driver is selected
self.win.toggle_edit_BTN.setEnabled(True)
self.win.delete_driver_BTN.setEnabled(True)
if solver.is_enabled:
self.win.toggle_edit_BTN.setText("EDIT")
else:
self.win.toggle_edit_BTN.setText("FINISH EDITING (ENABLE)")
self.win.add_pose_BTN.setEnabled(True)
self.win.add_driven_BTN.setEnabled(True)
self.win.select_driver_BTN.setEnabled(True)
self.win.copy_driven_trs_BTN.setEnabled(True)
self.win.paste_driven_trs_BTN.setEnabled(True)
selected_poses = self.get_selected_poses()
if selected_poses:
self.win.edit_pose_BTN.setEnabled(True)
self.win.delete_pose_BTN.setEnabled(True)
self.win.mirror_pose_BTN.setEnabled(True)
else:
self.win.add_pose_BTN.setEnabled(False)
self.win.add_driven_BTN.setEnabled(False)
self.win.toggle_edit_BTN.setEnabled(False)
self.win.delete_driver_BTN.setEnabled(False)
self.win.driver_transform_LINE.setText("")
self.win.select_driver_BTN.setEnabled(False)
# pose buttons
self.win.edit_pose_BTN.setEnabled(False)
self.win.delete_pose_BTN.setEnabled(False)
self.win.mirror_pose_BTN.setEnabled(False)
self.win.copy_driven_trs_BTN.setEnabled(False)
self.win.paste_driven_trs_BTN.setEnabled(False)
for i in range(self.win.driver_LIST.count()):
item = self.win.driver_LIST.item(i)
solver = item.data(QtCore.Qt.UserRole)
if not solver.is_enabled:
item.setText(solver.name + " (Editing)")
else:
item.setText(solver.name)
def pose_changed(self):
"""called when the pose selection is changed"""
selected_poses = self.get_selected_poses()
if selected_poses:
selected_pose = selected_poses[-1]
selected_solvers = self.get_selected_solvers()
solver = selected_solvers[-1]
if selected_pose in solver.pose_dict.keys():
solver.assume_pose(selected_pose)
else:
cmds.warning('Pose ' + selected_pose + ' not found in pose dictionary')
self.refresh_ui_state()
def load_drivers(self, selected=None):
"""loads all the RBF node drivers into the driver list widget"""
self.win.driver_LIST.clear()
drivers = [node for node in cmds.ls(type=self._node_name)]
selected_item = None
for driver in drivers:
item = QtWidgets.QListWidgetItem(driver)
solver = poseWrangler.UE4PoseDriver(existing_interpolator=driver)
item.setData(QtCore.Qt.UserRole, solver)
if selected and selected == driver:
selected_item = item
self.win.driver_LIST.addItem(item)
self.win.driver_LIST.sortItems(QtCore.Qt.AscendingOrder)
if selected_item:
self.win.driver_LIST.setCurrentItem(selected_item)
self.refresh_ui_state()
self.load_poses()
def load_poses(self):
"""loads the poses for the current driver"""
self.win.pose_LIST.clear()
selected_solvers = self.get_selected_solvers()
if selected_solvers:
solver = selected_solvers[-1]
for target in solver.pose_dict.keys() or []:
item = QtWidgets.QListWidgetItem()
item.setText(target)
item.setData(QtCore.Qt.UserRole, target)
self.win.pose_LIST.addItem(item)
# sort the items
self.win.pose_LIST.sortItems(QtCore.Qt.AscendingOrder)
def mirror_pose(self):
"""mirrors the selected pose for the current driver"""
selected_solvers = self.get_selected_solvers()
if not selected_solvers:
return
solver = selected_solvers[-1]
selected_poses = self.get_selected_poses()
if selected_poses:
for selected_pose in selected_poses:
pose_dict = solver.pose_dict
if selected_pose in pose_dict.keys():
poseWrangler.mirror_pose_driver(
solver.name, pose=selected_pose
)
else:
cmds.warning('Pose ' + selected_pose + ' not found in pose dictionary')
def mirror_driver(self):
"""mirrors all poses for the selected drivers"""
selected_solvers = self.get_selected_solvers()
for solver in selected_solvers:
poseWrangler.mirror_pose_driver(solver.name)
self.load_drivers()
def import_drivers(self, file_path=""):
"""imports the drivers"""
path = file_path or QtWidgets.QFileDialog.getOpenFileName(
self, "Pose Wrangler Format",
"", "JSON (*.json)"
)[0]
if path == "":
return
if not os.path.isfile(path):
OpenMaya.MGlobal.displayError(path + " is not a valid file.")
return
poseWrangler.import_drivers(path)
self.load_drivers()
self.load_poses()
def export_driver(self):
"""exports the selected drivers"""
selected_solvers = self.get_selected_solvers()
solver_names = []
for solver in selected_solvers:
solver_names.append(solver.name)
self._export(solver_names)
def export_all(self):
"""exports all rbf nodes"""
nodes = cmds.ls(type=self._node_name)
if not nodes:
return
self._export(nodes)
def _export(self, drivers):
"""exports the drivers passed in"""
file_path = QtWidgets.QFileDialog.getSaveFileName(None, "Pose Wrangler Format", "", "*.json")[0]
if file_path == "":
return
# do the export
poseWrangler.export_drivers(drivers, file_path)
def create_driver(self):
"""creates a new driver"""
sel = cmds.ls(sl=1)
if not sel:
OpenMaya.MGlobal.displayError(
'PoseWrangler: You must select a driving transform to CREATE a pose interpolator'
)
return
# takes all driven and the driving input is last
interp_name, ok = QtWidgets.QInputDialog.getText(self, 'text', 'Driver Name:')
if interp_name:
driver = poseWrangler.UE4PoseDriver()
driver.create_pose_driver_system(interp_name, sel[-1], sel[0:-1])
self._current_solver = driver
# refresh the combobox and set this interpolator as current
self.load_drivers(selected=self._current_solver.name if self._current_solver else None)
def delete_driver(self):
"""deletes the selected drivers"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.delete()
self.load_drivers()
self.load_poses()
def copy_driven_trs(self):
"""copies the driven TRS for pasting in different poses"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.copy_driven_trs()
def paste_driven_trs(self):
"""pastes the driven TRS and lets you multiply it"""
mult = self.paste_mult_DSPN.value()
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.paste_driven_trs(mult=mult)
def zero_base_poses(self):
"""zeros out the selected solver base poses"""
selected_solvers = self.get_selected_solvers()
if selected_solvers:
for solver in selected_solvers:
solver.zero_base_pose()
def _upgrade_scene(self):
file_path = upgrade.upgrade_scene(clear_scene=True)
LOG.info("Successfully Exported Current Scene")
self.event_upgrade_dispatch.upgrade.emit(file_path)
self.close()
def showUI():
"""show the UI"""
pose_wrangler_widget = PoseWrangler()
# show the UI
pose_wrangler_widget.show(dockable=True)