Technical Art Manager
'''
Save and Load Weights
Katelyn Malmsten, 2017
Used for saving and loading weights. Contains a UI and various functions to save weights
from objects, save them to JSON files, and load back onto objects.
'''
import json
import os
# This supports both PySide (up to Maya 2016) and PySide2 for Maya 2017
try:
from PySide2 import QtGui, QtCore, QtWidgets
except ImportError:
from PySide import QtGui, QtCore
import PySide.QtGui as QtWidgets
import pymel.core as pm
import maya.api.OpenMaya
import maya.api.OpenMayaAnim
WEIGHT_SAVE_DIR_DEFAULT = os.environ['temp']
QUICK_WEIGHT_SAVE_FILE = 'weightData_Temp'
WEIGHT_FILE_PREFIX = 'weightData'
# Global used for the UI so that it does not go out of scope
initUI = None
# ENTRY POINT
#================================================================================
def saveLoadWeights():
'''
Entry point to create the save/load UI
'''
global initUI
initUI = SaveLoadWeightsUI()
# UI
#================================================================================
class SaveLoadWeightsUI(QtWidgets.QMainWindow):
'''
UI for saving and loading weights
Allows user to define a directory to save/load to, choose from the saved files in that
location, and load to a mesh in the open scene
'''
def __init__(self, parent = None):
# check for existing windows, close if open
super(SaveLoadWeightsUI, self).__init__(parent)
self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)
self._initUI()
self.show()
def _initUI(self):
self.setWindowTitle('Save/Load Weights')
self.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding,
QtWidgets.QSizePolicy.MinimumExpanding)
self.setMinimumWidth(200)
self.setMinimumHeight(100)
mainWidget = QtWidgets.QFrame()
mainWidget.setObjectName('mainWidget')
self.setCentralWidget(mainWidget)
mainLayout = QtWidgets.QVBoxLayout(mainWidget)
mainLayout.setObjectName('mainLayout')
mainLayout.setAlignment(QtCore.Qt.AlignLeft|QtCore.Qt.AlignTop)
mainLayout.setContentsMargins(10, 10, 10, 10)
mainLayout.setSpacing(5)
browseLabel = QtWidgets.QLabel('Weights directory:')
#File browsing layout
browseLayout = QtWidgets.QHBoxLayout()
browseLayout.setObjectName('browseLayout')
browseLayout.setAlignment(QtCore.Qt.AlignLeft|QtCore.Qt.AlignTop)
browseLayout.setContentsMargins(0, 0, 0, 0)
browseLayout.setSpacing(5)
browseLayout.addWidget(browseLabel)
self.saveFilePath = QtWidgets.QLineEdit()
self.saveFilePath.setText(WEIGHT_SAVE_DIR_DEFAULT)
browseLayout.addWidget(self.saveFilePath)
self.browseButton = QtWidgets.QPushButton('Browse')
browseLayout.addWidget(self.browseButton)
#Save/Load buttons layout
row = QtWidgets.QHBoxLayout()
saveButton = QtWidgets.QPushButton('Save Weights from Selected Meshes')
loadButton = QtWidgets.QPushButton('Load Weight Data to Object')
row.addWidget(saveButton)
row.addWidget(loadButton)
#Second row of buttons
row2 = QtWidgets.QHBoxLayout()
saveSceneButton = QtWidgets.QPushButton('Save Weights from Scene')
refreshButton = QtWidgets.QPushButton('Refresh')
row2.addWidget(saveSceneButton)
row2.addWidget(refreshButton)
#Labels for file and mesh selection lists
columnLabels = QtWidgets.QHBoxLayout()
columnLabels.setAlignment(QtCore.Qt.AlignCenter|QtCore.Qt.AlignTop)
columnLabels.setContentsMargins(0, 0, 0, 0)
columnLabels.setSpacing(100)
weightFiles = QtWidgets.QLabel('Weight Data Files')
sceneMeshes = QtWidgets.QLabel('Meshes in Scene')
columnLabels.addWidget(weightFiles)
columnLabels.addWidget(sceneMeshes)
# Adding the weight manager to display weight files and scene meshes
self.weightManager = WeightFilesManager(self)
self.weightManager.populateWeightFiles(self.saveFilePath.text())
self.weightManager.populateSceneMeshes()
#Adding layouts to the main layout
mainLayout.addLayout(browseLayout)
mainLayout.addLayout(columnLabels)
mainLayout.addLayout(self.weightManager)
mainLayout.addLayout(row)
mainLayout.addLayout(row2)
self.setLayout(mainLayout)
#Connect Buttons
saveButton.clicked.connect(self.saveButtonOnClick)
loadButton.clicked.connect(self.loadButtonOnClick)
saveSceneButton.clicked.connect(self.saveSceneButtonOnClick)
refreshButton.clicked.connect(self.refreshLists)
self.browseButton.clicked.connect(self.browseOnClick)
def browseOnClick(self):
fileDir = QtWidgets.QFileDialog.getExistingDirectory()
if fileDir:
self.saveFilePath.setText(fileDir)
self.weightManager.populateWeightFiles(fileDir)
def refreshLists(self):
self.weightManager.populateSceneMeshes()
self.weightManager.populateWeightFiles(self.saveFilePath.text())
def saveButtonOnClick(self):
# Save weights from selected meshes
saveSelectedWeights(self.saveFilePath.text())
self.refreshLists()
def saveSceneButtonOnClick(self):
meshes = pm.filterExpand(pm.ls(tr=True), sm=12)
for mesh in meshes:
saveWeights(mesh, self.saveFilePath.text())
self.refreshLists()
def loadButtonOnClick(self):
weightFile = self.weightManager.getSelectedWeightFile().text()
weightDir = self.saveFilePath.text()
mesh = self.weightManager.getSelectedMesh()
if mesh and weightFile and weightDir:
loadWeights(pm.ls(mesh.text(), long=True)[0], weightDir, weightFile)
else:
pm.error('Must have a weight file, a mesh, and a directory selected to load weights')
# Text areas that list the weightfiles in the selected directory and meshes in scene
class WeightFilesManager(QtWidgets.QHBoxLayout):
'''
Contains the text list fields for the weight data files and the scene meshes as well
as the functions to access their data
'''
def __init__(self, parent):
super(WeightFilesManager, self).__init__(parent)
self._initUI()
def _initUI(self):
self.setObjectName('WeightFileManager')
self.setAlignment(QtCore.Qt.AlignLeft|QtCore.Qt.AlignTop)
self.setContentsMargins(5, 5, 5, 5)
self.setSpacing(5)
self.weightFilesList = QtWidgets.QListWidget()
self.sceneMeshes = QtWidgets.QListWidget()
self.addWidget(self.weightFilesList)
self.addWidget(self.sceneMeshes)
def populateWeightFiles(self, fileDir):
#Empty the list of weight files, then populate it
self.weightFilesList.clear()
weightDataFiles = []
if os.path.exists(fileDir):
for f in os.listdir(fileDir):
if f.endswith(".json") and f.startswith(WEIGHT_FILE_PREFIX):
weightDataFiles.append(f)
self.weightFilesList.addItems(weightDataFiles)
else:
pm.error('{0} is not a valid path'.format(fileDir))
def populateSceneMeshes(self):
self.sceneMeshes.clear()
meshes = pm.filterExpand(pm.ls(tr=True), sm=12)
for mesh in meshes:
self.sceneMeshes.addItem(str(mesh))
def getSelectedWeightFile(self):
return self.weightFilesList.currentItem()
def getSelectedMesh(self):
return self.sceneMeshes.currentItem()
# WEIGHT GET AND SET
#================================================================================
def getWeights(mesh):
'''
Gets the weightdata from the specified mesh
:param mesh: Mesh to get the weights from
:return: dictionary containing the influences and weights
'''
#Related skin cluster
skincluster = pm.mel.findRelatedSkinCluster(mesh)
if skincluster:
node = getMObject(skincluster)
skinClusterFn = maya.api.OpenMayaAnim.MFnSkinCluster(node)
influenceObjectPaths = skinClusterFn.influenceObjects()
influencePaths = {}
influences = []
for influence in xrange(len(influenceObjectPaths)):
influencePath = influenceObjectPaths[influence].fullPathName()
influencePaths[int(skinClusterFn.indexForInfluenceObject(influenceObjectPaths[influence]))] = influence
influences.append(influencePath)
# Get the weightlist and weights attributes
weightListPlug = skinClusterFn.findPlug('weightList', False)
weightsPlug = skinClusterFn.findPlug('weights', False)
weightListAttr = weightListPlug.attribute()
weightsAttr = weightsPlug.attribute()
# Store the weights
# The stucture is a dictionary with the vertID as the key, value of another dictionary of { influence ID : weights for influence}
# Final format is {vertID : {influence ID : weights for influence}}
weights = {}
for vertex in xrange(weightListPlug.numElements()):
vertexWeights = {}
weightsPlug.selectAncestorLogicalIndex(vertex, weightListAttr)
# non-zero weights for the vertex
weightInfluenceIds = weightsPlug.getExistingArrayAttributeIndices()
# duplicate the current weightPlug
influencePlug = maya.api.OpenMaya.MPlug(weightsPlug)
#Filling vertexWeights dictionary with {influence : weights}
for influenceId in weightInfluenceIds:
influencePlug.selectAncestorLogicalIndex(influenceId, weightsAttr)
try:
vertexWeights[influencePaths[influenceId]] = influencePlug.asDouble()
except KeyError:
# if the influence has been removed
pass
weights[vertex] = vertexWeights
return {'influences' : influences, 'weights' : weights}
else:
pm.error('No skin cluster on {0}'.format(mesh))
#--------------------------------------------------------------------------------
def setWeights(mesh, influences, weights):
'''
Loads the weights from weightData onto the selected unskinned mesh
:param mesh: The mesh to load the weights onto
:param influences: list of influences
:param weights: Dictionary containing vertexes and weights
'''
if pm.mel.findRelatedSkinCluster(mesh):
raise ValueError('setWeights error: Passed in mesh already has a skin. Unbind the skin cluster first')
# Create skin cluster
skinCluster = pm.skinCluster(influences +[mesh.name()], toSelectedBones = True)
node = getMObject(skinCluster.name())
skinClusterFn = maya.api.OpenMayaAnim.MFnSkinCluster(node)
influenceObjectPaths = skinClusterFn.influenceObjects()
influencePaths = {}
for influence in xrange(len(influenceObjectPaths)):
influencePath = influenceObjectPaths[influence].fullPathName()
influencePaths[int(skinClusterFn.indexForInfluenceObject(influenceObjectPaths[influence]))] = influencePath
# Clear weighting to create a clean base to rebuild weighting onto
# Unlock influences
for influence in influences:
pm.setAttr('{0}.liw'.format(influence))
# prune weights and then restore normalize settings
skinNorm = pm.getAttr(skinCluster.normalizeWeights)
if skinNorm != 0:
pm.setAttr(skinCluster.normalizeWeights, 0)
pm.skinPercent(skinCluster, mesh, pruneWeights =100, normalize = False)
if skinNorm != 0:
#skinCluster.normalizeWeights.set(skinNorm)
pm.setAttr(skinCluster.normalizeWeights, skinNorm)
# Set weights
for vertex, weightData in weights.items():
weightListAttr = skinCluster.weightList[vertex]
for influenceId, weight in weightData.items():
pm.setAttr(weightListAttr.weights[influenceId], weight)
#--------------------------------------------------------------------------------
def weightDataHandler(action, fileDir, fileName = QUICK_WEIGHT_SAVE_FILE, mesh=None):
'''
Handles reading or writing data to a json file
:param string action: The type of operation, r = read weights from file, w = save weights
:param string fileDir: The directory path containing weight files
:param string fileName: Name of the weight file to save to or load from
:param mesh: The source mesh (for use in saving weights)
'''
if not os.path.exists(fileDir):
raise ValueError("The selected directory for saving/loading weights does not exist")
weightFilePath = os.path.join(fileDir, fileName)
if action != 'w' and action != 'r':
raise ValueError('Action must be either "r" or "w"')
try:
with open(weightFilePath, action) as weightFile:
if action == 'r':
weightData = json.load(weightFile)
return weightData
elif mesh:
weightData = getWeights(mesh)
json.dump(weightData, weightFile)
else:
raise ValueError('weightDataHandler error: A valid mesh must be passed in to save weights from')
except Exception as exc:
raise RuntimeError('Failed to read or write weight data for {0}. Exception details: {1}'.format(weightFilePath, exc.args[0]))
# SAVING AND LOADING WEIGHTS
#================================================================================
def saveWeights(mesh, fileDir):
'''
Save the weighting of the currently selected mesh.
If multiple are selected, will use the first selected object
:param mesh: mesh to save weights from
:param string fileDir: The directory path for locating weights
'''
# Validating input
validSaveWeightsTarget = True
skinCluster = pm.mel.findRelatedSkinCluster(mesh)
if not pm.listRelatives(mesh, shapes = True):
print '{0} does not have a mesh'.format(mesh)
validSaveWeightsTarget = False
if not skinCluster:
print '{0} does not have a skin cluster'.format(mesh)
validSaveWeightsTarget = False
if validSaveWeightsTarget:
# Write the json file
fileName = '{0}_{1}.json'.format(WEIGHT_FILE_PREFIX, mesh)
weightDataHandler('w', fileDir, fileName, mesh)
print 'Success: Weights saved from {0}'.format(mesh)
else:
# Will not save weights if there are no weights to save
# returns rather than raising an exception so that doesn't prevent saving from additional meshes when running through a list
return
#--------------------------------------------------------------------------------
def loadWeights(mesh, fileDir, fileName = QUICK_WEIGHT_SAVE_FILE):
'''
Loads weighting data from weight file onto selected mesh.
:param mesh: mesh to load the weights on to
:param string fileDir: The directory path for locating weights
:param string fileName: Name of the source weights file
'''
#Check for a valid selection
if len(pm.listRelatives(mesh, shapes = True)) > 0:
#Clean off an existing skin cluster before loading new weights
skinCluster = pm.mel.findRelatedSkinCluster(mesh)
if skinCluster:
pm.skinCluster(skinCluster, edit = True, unbind = True)
else:
print '{0} does not have a mesh'.format(mesh)
return
# get weightdata from json file
print 'Retrieving weight data'
weightData = weightDataHandler('r', fileDir, fileName, mesh)
# Apply weights to mesh
setWeights(mesh, weightData['influences'], weightData['weights'])
print 'Success: Weights set on {0}'.format(mesh)
#--------------------------------------------------------------------------------
def saveSelectedWeights(fileDir):
'''
Saves the weights of selected meshes
:param fileDir: The directory path for weight data
'''
meshes = getSelectedMeshes()
for mesh in meshes:
saveWeights(mesh, fileDir)
# HELPER FUNCTIONS
#================================================================================
def getMObject(name):
'''
Gets the MObject for the passed in object
:param name: the object to retrieve an mObject for
:return node: MObject
'''
selectionList = maya.api.OpenMaya.MSelectionList()
selectionList.add(name)
node = selectionList.getDependNode(0)
return node
#--------------------------------------------------------------------------------
def getSelectedMeshes():
'''
Gets list of selected objects that have a mesh
'''
selected = pm.selected()
if not selected:
raise ValueError('No object is selected, please select a mesh')
for item in selected:
# Anything without a mesh should not be left in this list
if not pm.listRelatives(item, shapes = True):
selected.remove(item)
return selected