Maya averageRBGA

maya_averageRGBA.py

#!/usr/bin/env python
#coding:utf-8
'''Loops selection to average RGBA on vertices.'''

__author__ = 'hogjonny'

import maya
from maya import cmds,mel

global gDebug
gDebug = 0

global gMainProgressBar
gMainProgressBar = mel.eval('$tmp = $gMainProgressBar')
#----------------------------------------------------------------------


#----------------------------------------------------------------------
class UndoContext(object):
    """
    This is a Class, it creates a undo context chunk
    """
    def __enter__(self):
        cmds.undoInfo(openChunk=True)
    def __exit__(self, *exc_info):
        cmds.undoInfo(closeChunk=True)
#----------------------------------------------------------------------


#----------------------------------------------------------------------
# helper
def chunks(l, n):
    """ Yield successive n-sized chunks from l."""
    for i in xrange(0, len(l), n):
        yield l[i:i+n]
#----------------------------------------------------------------------


#----------------------------------------------------------------------
def get_mobject( node):
    selectionList = OpenMaya.MSelectionList()
    selectionList.add(node)
    oNode = OpenMaya.MObject()
    selectionList.getDependNode(0, oNode)
    return oNode
#----------------------------------------------------------------------


#----------------------------------------------------------------------
# main looper
def bp_avgVertColorRGBA(vertexColorList, vselChunk, selmask, vertloop, vSelTotal, chunkSize, gDebug=gDebug):
    
    for v in vselChunk:
        if cmds.progressBar(gMainProgressBar, query=True, isCancelled=True ):
                break
        
        if gDebug:
            print 'vertex: {0}'.format(v)
            
        # maya api color
        #omRgba = OpenMaya.MColor(0,0,0,0)

        # maya /python roll your own ... index a list
        rgba = [0,0,0,0]

        if gDebug:
            print 'initializing color: {0}'.format(rgba)

        # get and average all the vertex face components
        vtxfaces = set( cmds.ls( cmds.polyListComponentConversion( v, fv=1, tvf=1 ), fl=1 ) )
        vtxfaces = vtxfaces & selmask
        av = len(vtxfaces)
        if gDebug:
            print 'num of vf: {0}'.format(av)

        for vf in vtxfaces:
            vfc = cmds.polyColorPerVertex(vf,r=1,g=1,b=1,a=1,q=1)
            if gDebug:
                print 'vertex face rgba is: {0}'.format(vfc)
            rgba[0] += (vfc[0] / av)
            rgba[1] += (vfc[1] / av)
            rgba[2] += (vfc[2] / av)
            rgba[3] += (vfc[3] / av)

        if gDebug:
            print 'computed color is: {0}'.format(rgba)
            
        # pack an api color tuple
        omRgba = OpenMaya.MColor(rgba[0],rgba[1],rgba[2],rgba[3])
            
        # put it in the vertexColorList
        vertexColorList.set(omRgba, vertloop)

        #cmds.select(v, replace=True)
        #try:
            #cmds.polyColorPerVertex( r=rgba[0], g=rgba[1], b=rgba[2], a=rgba[3],
                                     #colorDisplayOption=True, notUndoable=True )
        #except:
            #print 'polyColorPerVertex fialed, vert: {0}'.format(v)
            #pass
        
        cmds.progressBar(gMainProgressBar, edit=True, step=1,
                         status='bp_avgVertColorRGBA() >> number of verts: {0}'.format(vSelTotal-vertloop))
        vertloop += 1
        
    return vertloop
#----------------------------------------------------------------------


#----------------------------------------------------------------------
def main():
    # standard maya/python
    sel = cmds.ls(sl=1,fl=1)
    if gDebug:
        print 'selection: {0}'.format(sel)
        
    # maya api
    for obj in sel:
        # on object at a time
        # get at the objects node for the api calls
        objApi = cmds.listRelatives(obj, s=True,ni=True)[0]
        selObjApi = get_mobject(objApi)
        meshFnApi = OpenMaya.MFnMesh(selObjApi)
        
        # initialize a list, for storing computed rgba values,
        # dump the current meshes values into it ... we will replace them
        # then dump this muted list back onto the object
        vertexColorList = OpenMaya.MColorArray()
        meshFnApi.getVertexColors(vertexColorList)
        
        selmask = set(cmds.ls(cmds.polyListComponentConversion(obj, tvf=1),fl=1))
        if gDebug:
            print 'selmask: {0}'.format(selmask)
        
        vsel = cmds.ls(cmds.polyListComponentConversion(obj, tv=1),fl=1)
        vSelTotal = len(vsel)
        if gDebug:
            print 'vsel num: {0}'.format(len(vsel))
        
        # chunk the list (because if it's huge the loop will just crash maya runtime)
        chunkSize = 512
        chunkList = chunks(vsel, chunkSize)
            
        cmds.progressBar( gMainProgressBar,
                          edit=True,
                          beginProgress=True,
                          isInterruptable=True,
                          status='bp_avgVertColorRGBA() >> number of verts total: {0}'.format(len(vsel)),
                          maxValue=len(vsel),
                          progress=0 )
        
        if cmds.progressBar(gMainProgressBar, query=True, isCancelled=True ):
                break
        
        global vertloop
        vertloop = 0
        
        # loop through chunks
        for index, chunk in enumerate(chunkList):
            if gDebug:
                print 'Working on Chunk: {0}'.format(index)
            with bpUndo.UndoContext():
                print 'The vertloop is: {0}'.format(vertloop)
                vertloop = bp_avgVertColorRGBA(vertexColorList, chunk, selmask, vertloop, vSelTotal, chunkSize)
        
        # now that we have stored off all of the averaged values in the vertexColorList
        # we can push them back onto the object via the api in one call
        fnComponent = OpenMaya.MFnSingleIndexedComponent()
        fullComponent = fnComponent.create( OpenMaya.MFn.kMeshVertComponent )
        fnComponent.setCompleteData( vertexColorList.length() )
        vertexIndexList = OpenMaya.MIntArray()
        fnComponent.getElements(vertexIndexList)
        # set the colors
        meshFnApi.setVertexColors(vertexColorList, vertexIndexList, None)
    
    cmds.progressBar(gMainProgressBar, edit=True, endProgress=True, status='bp_avgVertColorRGBA(): Finished all verts!')
    cmds.select(sel, replace=True)
#----------------------------------------------------------------------


#----------------------------------------------------------------------
# main loop
if __name__ == "__main__":
    main()
#----------------------------------------------------------------------