Files
MetaBox/Scripts/Animation/epic_pose_wrangler/v1/poseWranglerUI.py
2025-04-17 04:52:48 +08:00

659 lines
23 KiB
Python

# Copyright Epic Games, Inc. All Rights Reserved.
from Qt import QtWidgets
from Qt import QtCore
# from Qt import QtUiTools
from Qt.QtCompat import loadUi
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)
self.win = loadUi(file_path)
finally:
# Always close the UI file regardless of loader result
# ui_file.close()
pass
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)