yeti cache export script for maya

cacheExport.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
import sys, pprint
from pysideuic import compileUi
pyfile = open("/nas/projects/development/pipeline/bs_pipeline/cacheExportUI.py", 'w')
compileUi("/nas/projects/development/pipeline/bs_pipeline/ui/cacheExport.ui", pyfile, False, 4, False)
pyfile.close()
'''

from maya import OpenMayaUI as omui
from pymel.core import *
from PySide import QtCore, QtGui
from shiboken import wrapInstance
from bs_pipeline.cacheExportUI import Ui_cacheExport
import os
import re
import addOffsetAttributes
reload(addOffsetAttributes)


# Global Variables Initialization
WINDOW_TITLE = 'Write Cache'
WINDOW_VERTION = "1.8.5"
WINDOW_NAME = 'cacheExport'

def mayaMainWindow():
    try:
        mayaWinPtr = omui.MQtUtil.mainWindow()
        mayaWin = wrapInstance(long(mayaWinPtr), QtGui.QWidget)
        return mayaWin
    except:
        return None

class cacheExportUI(QtGui.QDialog):
    def __init__(self, parent = mayaMainWindow(), geo=None, cacheType=None, cachePath=None, cacheStart=None, cacheEnd=None, cacheAtOrig=True, cacheOffsetFile=None):
        super(cacheExportUI, self).__init__(parent)
        self.ui = Ui_cacheExport()
        self.ui.setupUi(self)
        
        self.setWindowTitle('{0} {1}'.format(WINDOW_TITLE, str(WINDOW_VERTION)))
            
        QtCore.QObject.connect(self.ui.cacheBtn, QtCore.SIGNAL("clicked()"), self.expCache)
        QtCore.QObject.connect(self.ui.browseBtn, QtCore.SIGNAL("clicked()"),self.browse)
        #QtCore.QObject.connect(self.ui.renderableChk, QtCore.SIGNAL("toggled(bool)"), self.renderableOnly)
        QtCore.QObject.connect(self.ui.helpBtn, QtCore.SIGNAL("clicked()"), self.help)
        QtCore.QObject.connect(self.ui.startSpin, QtCore.SIGNAL("valueChanged(int)"), self.setStart)
        QtCore.QObject.connect(self.ui.endSpin, QtCore.SIGNAL("valueChanged(int)"), self.setEnd)
        QtCore.QObject.connect(self.ui.cachePathLine, QtCore.SIGNAL("textChanged(QString)"), self.setPath)
        QtCore.QObject.connect(self.ui.cacheTypeBox, QtCore.SIGNAL("currentIndexChanged(QString)"), self.typeChange)
        QtCore.QObject.connect(self.ui.cacheAtOrigChk, QtCore.SIGNAL("toggled(bool)"), self.origCacheTgl)

        self.geo = False
        self.cacheFile = False
        self.simType = self.ui.cacheTypeBox.currentText()
        self.min = playbackOptions(q=True, minTime=True)
        self.max = playbackOptions(q=True, maxTime=True)
        
    
        self.ui.startSpin.setValue(self.min)
        self.ui.endSpin.setValue(self.max)
        self.cacheAtOrig = cacheAtOrig
        self.disableConfirmation = False
        self.cacheOffsetFile = cacheOffsetFile
        
    def origCacheTgl(self):
        self.cacheAtOrig = self.ui.cacheAtOrigChk.isChecked()
        
            
    def typeChange(self):
        self.simType = self.ui.cacheTypeBox.currentText()

        if self.simType =='Rig' or self.simType=='Cam':
            self.ui.cacheAtOrigChk.setEnabled(1)
        else:
            self.ui.cacheAtOrigChk.setEnabled(0)
            
        
    def setStart(self):
        self.min = self.ui.startSpin.value()
    
    
    def setEnd(self):
        self.max = self.ui.endSpin.value()
    
    
    def setPath(self):
        self.cacheFile = self.ui.cachePathLine.text()
    
    
    
    def browse(self):
        if not self.ui.cacheTypeBox.currentText() == 'Yeti Hair/Fur':
            abcFilter = 'bs_alembic abc (*.abc);;All Files (*.*)'
            abcF = fileDialog2(cap='Export to alembic...', fileFilter=abcFilter, dialogStyle=2)
            if abcF:
                abcF = os.path.splitext(abcF[0])[0]
                self.cacheFile = '%s.abc'%abcF
                self.ui.cachePathLine.setText(self.cacheFile)
            else:
                return False
        else:
            cacheDir = fileDialog2(cap='Select Cache Root Folder...', fm = 2, okc="Set Root", dialogStyle=2)
            self.ui.cachePathLine.setText(cacheDir[0])
    '''            
    def renderableOnly(self):
        if self.ui.renderableChk.isChecked():
            self.ui.cacheTypeBox.setEnabled(0)
            self.ui.label.setEnabled(0)
        else:
            self.ui.cacheTypeBox.setEnabled(1)
            self.ui.label.setEnabled(1)
    
    '''
    
    def help(self):
        os.system('acroread /nas/projects/development/productionTools/bs_pipeline/doc/cacheExport-v%s.pdf & '%WINDOW_VERTION)
        print("cacheExport-v%s:\n1. Select the Asset Root group\n2. select the Cache Type.\n3. Browse where your cache should go.\n4. Cache it."%WINDOW_VERTION)
    
    def expCache(self):
        cacheExport(geo=self.geo, cacheType=self.simType, cachePath=self.cacheFile, cacheStart=self.min, cacheEnd=self.max, cacheAtOrig=self.cacheAtOrig, cacheOffsetFile=self.cacheOffsetFile).expCache()
        
        
class cacheExport:
    def __init__(self, geo=None, cacheType=None, cachePath=None, cacheStart=None, cacheEnd=None, cacheAtOrig=True, cacheOffsetFile=None):
        '''
        geo\tThe asset root group you for the asset you want to cache
        cacheType\tString\t"Rig","Cam", "Pfx Hair", "Yeti Hair/Fur", "Cloth", "Fluid", "Particles", "Heavy Geo"
        cacheStart
        cacheEnd
        cachePath this should be a folder path for Yeti cache, .ass for heavy Geo or .abc for anything else
        cacheAtOrig: Caches assets of type 'Rig', 'Prop' or 'Cam' near to the origin (generating a text file for offset if no one passed in here)
        cacheOffsetFile: the cache offset file will be written/read when caching near the origin.
        '''

        self.geo = geo
        self.cacheFile = cachePath
        self.simType = cacheType
        self.min = cacheStart
        self.max = cacheEnd
        self.cacheAtOrig = cacheAtOrig
        self.disableConfirmation = False
        self.cacheOffsetFile = cacheOffsetFile
        
        if objExists('originTranslateLocator'):
            self.disableConfirmation = True
            
        
        '''
        if geo is not None:
            self.disableConfirmation = True
            if cacheOffsetFile is None and cacheAtOrig:
                error('An offset file "cacheOffsetFile" needs to be passed when running in batch mode')
        '''    
    
    def expCache(self):
        print 'cacheExport-v%s'%WINDOW_VERTION
        ignoreTags = False
                    
        if self.geo:
            select(self.geo, r=True)
            
        roots = ls(sl=True)
            
        simTagged=[]
        simType = self.simType
        startFrame=self.min
        endFrame=self.max
        
        if simType=='Yeti Hair/Fur':
            chkSel = ls(sl=True)
            if chkSel:
                furGrps = ls(sl=True, type='transform')
                if not furGrps:
                    furGrps = ls(sl=True, type='shape')
                    
            #self.writeYetiCache(furGrps, startFrame, endFrame, self.ui.cachePathLine.text())
            self.writeYetiCache(furGrps, startFrame, endFrame, self.cacheFile)
            return
        elif simType=='Pfx Hair':
            simTag = 'simHair'
        elif simType=='Cloth':
            simTag = 'simCloth'
        elif simType =='Rig':
            simTag = 'rig'
        elif simType =='Cam':
            ignoreTags = True
        else:
            warning('this feature is not supported yet!!')
            return
        # For all types except props
        if not ignoreTags:
            for root in roots:
                meshes = ls(root, dag=True, type='transform')
                for mesh in meshes:
                    if attributeQuery('sim', node=mesh, ex=True):
                        taggedValue = getAttr('%s.sim'%mesh)
                        if taggedValue.split('_')[0] == simTag:
                            simValue = [taggedValue.split('_')[1], mesh]
                            simTagged.append(simValue)
                    
                    if attributeQuery('level_of_detail', node=mesh, ex=True):
                        try:
                            setAttr('%s.level_of_detail'%mesh, 2)
                        except:
                            pass
                        try:
                            setAttr('%s.cloth'%mesh, 0)
                        except:
                            pass
                        try:
                            setAttr('%s.facial'%mesh, 0)
                        except:
                            pass
                        try:
                            setAttr('%s.ctrlsVis'%mesh, 0)
                        except:
                            pass
                        try:
                            setAttr('%s.mesh_display'%mesh, 0)
                        except:
                            pass
                        try:
                            setAttr('%s.Hair'%mesh, 0)
                        except:
                            pass
                        try:
                            setAttr('%s.hairVis'%mesh, 0)
                        except:
                            pass
                
            simObjects = sorted(simTagged)
                
            if simObjects:
                #abc = self.cacheTagged(simObjects, cacheFile=self.ui.cachePathLine.text())
                abc = self.cacheTagged(simObjects, cacheFile=self.cacheFile)
                
                if abc:
                    print('Export to %s Done.'%abc)
                else:
                    print('Export cancelled!')
            else:
                error('Nothing marked "sim" under the current selection group(s), please use "aSim" tool to mark sim objects.')
        else: # This is a prop export and it doesn't have sim tag
            simObjects = ls(sl=True)
            if len(simObjects) > 1:
                error('Only one prop can be cached at a time. Please select the prop asset root group and try again')
                return
            else:
                # abc = self.cacheTagged(simObjects, cacheFile=self.ui.cachePathLine.text())
                abc = self.cacheTagged(simObjects, cacheFile=self.cacheFile)
                if abc:
                    print('Export to %s Done.'%abc)
                else:
                    print('Export cancelled!')
        
    
    def cacheTagged(self, geo, cacheFile=''):
        '''
        write cache tagged objects with multiple roots.
        '''
        startFrame=self.min # self.ui.startSpin.value()
        endFrame=self.max # self.ui.endSpin.value()
        
        customAttrs=[]
        abcAttr=''
        
        # custom shape attrs to .geomArb
        rootChildren=ls(geo, dag=True, type='shape')
        for child in rootChildren:
            attrs = listAttr(child, ud=True)
            if attrs:
                for attr in attrs:
                    if attr not in customAttrs:
                        customAttrs.append(attr)
                        abcAttr = abcAttr + '-attr %s '%attr
        
        '''
        # custom trasformNodes attrs to .geomArb  
        rootChildrenXF=ls(geo, dag=True, type='transform')
        for child in rootChildrenXF:
            attrs = listAttr(child, ud=True)
            if attrs:
                for attr in attrs:
                    if attr not in customAttrs:
                        customAttrs.append(attr)
                        abcAttr = abcAttr + '-attr %s '%attr
        '''             
        if 'sim' not in customAttrs:
            abcAttr = abcAttr + '-attr sim '
        
        abcAttr = abcAttr + '-attr panZoomEnabled -attr horizontalPan -attr verticalPan -attr zoom -attr renderPanZoom '
        
        # Save the proper Rig position to file, then move it near the origin.
        if self.simType =='Rig' and self.cacheAtOrig:
            confirmed = self.cacheAtOrigConfirm()
            posFilename = '%s.txt'%os.path.splitext(cacheFile)[0]
            if confirmed=='Use existing':
                # This variable is the input cacheOffset file (self.cacheOffsetFile)
                addOffsetAttributes.autoExtractOffset(reset=False, filename=posFilename, ws=True, existent=True)
            elif confirmed == 'Generate New':
                # call the command that generate the text file
                addOffsetAttributes.autoExtractOffset(reset=False, filename=posFilename, ws=True, existent=False)
                pass
            else:
                return
                
        elif self.simType == 'Cam' and self.cacheAtOrig: # for props we should use the other function
            posFilename = '%s.txt'%os.path.splitext(cacheFile)[0]
                # This variable is the input cacheOffset file (self.cacheOffsetFile)
            addOffsetAttributes.autoExtractOffset(reset=False, filename=posFilename, ws=True, existent=True)
        else:
            pass
        
                
        rootArg = ''
        roots = ls(geo)
        astRoot = ls(sl=True)[0]
        sel = ls(roots, dag=True)
        select(sel)        
        
        for root in roots:
            rootArg = rootArg + '-root %s '%root
        if self.simType =='Pfx Hair':
            AbcExport(j="-frameRange %s %s %s -eulerFilter -writeVisibility -uvWrite -root %s -sl -file %s"%(startFrame, endFrame, abcAttr, astRoot, cacheFile))
        elif self.simType == 'Cam':
            AbcExport(j="-frameRange %s %s %s -worldSpace -df hdf -eulerFilter -writeVisibility -uvWrite -root %s -sl -file %s"%(startFrame, endFrame, abcAttr, astRoot, cacheFile))
        else:
            AbcExport(j="-frameRange %s %s %s -worldSpace -eulerFilter -writeVisibility -uvWrite -root %s -sl -file %s"%(startFrame, endFrame, abcAttr, astRoot, cacheFile))
            
        # Restore the character rig proper position:
        select(astRoot, r=True)
        ## addOffsetAttributes.autoExtractOffset(reset=True)
        
        return cacheFile

    def cacheAtOrigConfirm(self):
        if self.disableConfirmation:
            return 'Use existing'
        
        confirmed = confirmDialog(title='Confirm...', message='You chosed to cache near to the origin, this requires a pre-existing scene offset file..', button=['Generate New','Use existing', 'Cancel'])
        return confirmed
        
    
    def cacheRenderable(self, cacheFile):
        startFrame=self.min
        endFrame=self.max

        customAttrs=[]
        abcAttr=''
        root = ls(sl=True)[0]
        rootChildren = ls(root, dag=True, type='shape')
        for child in rootChildren:
            attrs = listAttr(child, ud=True)
            if attrs:
                for attr in attrs:
                    if attr not in customAttrs:
                        customAttrs.append(attr)
                        abcAttr = abcAttr + '-attr %s '%attr
                        
        if 'sim' not in customAttrs:
            abcAttr = abcAttr + '-attr sim '
            
        rootArg = "-root %s"%root
        
        AbcExport(j="-frameRange %s %s %s -worldSpace -renderableOnly -eulerFilter -uvWrite %s -file %s"%(startFrame, endFrame, abcAttr, rootArg, cacheFile))
        return cacheFile
    
    
    def writeYetiCache(self, furGrps, start, end, cacheRoot):
        sc = sceneName()
        tmp = os.path.splitext(os.path.basename(sc))[0].split('_')
            
        for root in furGrps:
            yetiShapes = []
            xforms = ls(root, dag=True, type='shape')
            for xform in xforms:
                if attributeQuery('cacheFileName', node=xform, ex=True):
                    yetiShapes.append(xform)
        
            if yetiShapes:
                try:
                    cacheDirname = '%s_%s_%s'%(tmp[0], tmp[1], tmp[2])
                except:
                    cacheDirname = os.path.splitext(os.path.basename(sc))[0]
                
                for fur in yetiShapes:
                    try:
                        os.makedirs(os.path.join(cacheRoot, cacheDirname, fur))
                    except:
                        pass
                    
                for fur in yetiShapes:
                    cacheFilePath = os.path.join(cacheRoot, cacheDirname, fur, (fur + ".%04d.fur" ) )		
                    setAttr('%s.fileMode'%fur, 0)
                    print cacheFilePath
                
                    pgYetiCommand(fur, writeCache=cacheFilePath, 
                                    range=(start, end),
                                    samples=3, 
                                    sampleTimes="-0.25 0 0.25",
                                    updateViewport=False) 
                    
                    setAttr('%s.cacheFileName'%fur, cacheFilePath)
                    setAttr('%s.fileMode'%fur, 2)
                    
            else:
                warning("Couldn't find any Yeti nodes to cache under '%s'!!"%root)


def setYetiCache():
    try:
        astRoot = ls(sl=True)[0]
        base_elem = astRoot.split('_')[0]
    except Exception as e:
        print e
        
    cachePath = fileDialog2(cap='Select Cache Root Folder...', fm = 2, okc="Set Root", dialogStyle=2)[0]
    elem_name = re.search((".*({0}[a-zA-Z0-9]*)").format(base_elem), cachePath, re.I).groups()[0]
    
    yetiShapes = []
    xforms = ls(astRoot, dag=True, type='shape')
    for xform in xforms:
        if attributeQuery('cacheFileName', node=xform, ex=True):
            yetiShapes.append(xform)

    for fur in yetiShapes:
        fur_elem = str(fur).split('|')[-1]
        if not os.path.exists(os.path.join(cachePath, fur_elem)):
            pat = re.compile(base_elem, re.I)
            if not os.path.exists(os.path.join(cachePath, pat.sub(elem_name, fur_elem))):
                elem_name = re.sub('([a-zA-Z])', lambda x:x.groups()[0].upper(), elem_name, 1)
            fur_elem = pat.sub(elem_name, fur_elem)
        cacheFilePath = os.path.join(cachePath, fur_elem, (fur_elem + ".%04d.fur" ) )
        print cacheFilePath

        setAttr('%s.fileMode'%fur, 0)
        setAttr('%s.cacheFileName'%fur, cacheFilePath, type="string")
        setAttr('%s.fileMode'%fur, 2)


    
def show():
	if cmds.window(WINDOW_NAME, exists=True, q=True):
		cmds.deleteUI(WINDOW_NAME)
		dialog = None
	
	dialog = cacheExportUI()
	dialog.show()
	return dialog