Help to process high mesh normal on low mesh~ for Maya

normalWithVtxColor.py

# coding=utf-8

# Script by MineClever, help to process high mesh normal on low mesh~

import maya.api.OpenMaya as api
import maya.cmds as cmds

class MSingleton(type):
    """metaclass for set singleton class"""
    __debug = True
    
    def __init__(self, *args, **kwargs):
        self.__instance = None
        super(MSingleton,self).__init__(*args, **kwargs)

    def __call__(self, *args, **kwargs):
        if self.__instance is None:
            if self.__debug: print(u"current class {} create singleton class".format(self.__class__))
            self.__instance = super(MSingleton,self).__call__(*args, **kwargs)
        return self.__instance

class MSelectionHelper(object):
    __metaclass__ = MSingleton
    
    __debug = False
    tempSelectionList = api.MSelectionList()

    @classmethod
    def getCurMeshSelections(cls):
        """Return all mesh shape current selected"""
        selList = api.MGlobal.getActiveSelectionList() # type: api.MSelectionList
        shapeList = cls.tempSelectionList
        shapeList.clear()
        numOfSelected = selList.length() # type: int
        if cls.__debug: print("Current selection count : {}".format(numOfSelected))
        
        # NOTE: check if empty list
        if not (numOfSelected >0) :
            return shapeList

        # NOTE: add all shape into selection list
        for i in range(numOfSelected):
            dagPath = selList.getDagPath(i) # type: api.MDagPath
            if (dagPath.numberOfShapesDirectlyBelow() > 0): dagPath.extendToShape()
            # NOTE: check if mesh type
            if cls.__debug: print("selection[{}]:\"{}\" @ type: {}".format(i, dagPath, dagPath.apiType()))
            if (dagPath.apiType() == 296) : 
                shapeList.add(dagPath)
        if cls.__debug:  print(u"shape list count: {}".format(shapeList.length()))
        return shapeList

class MMProcessorBar(object):
    __debug = False
    __metaclass__ = MSingleton

    __processorCount = 0
    __processorTotalCount = 0
    __bShowProcessBar = False

    vtxUpdateThresholdNum = 40960
    
    @classmethod
    def setProcessorCountBarTotalNum(cls,inTotalCount):
        # type: (int) -> None
        if not cls.__bShowProcessBar:
            if cls.__debug: print("Is not able to set a processor bar now")
            return
        
        cls.__processorTotalCount = inTotalCount
        cmds.progressWindow( 
        title='Process Vertex Data',
        progress=0,
        status=("Done {0}/{1}\t".format(0, inTotalCount)),
        isInterruptable=True,
        min=0,
        max=inTotalCount)

    @classmethod
    def updateProcessorCountBar(cls,inCurCount):
        # type: (int) -> None
        if not cls.__bShowProcessBar:
            return
        
        cls.__processorCount = inCurCount
        if cls.__bShowProcessBar :
            if (cmds.progressWindow( query=True, progress=True ) >= cls.__processorTotalCount) or \
            (cmds.progressWindow( query=True, isCancelled=True ) ) :
                cls.endProcessorCountBar()
            else:
                cmds.progressWindow( edit=True, progress=inCurCount, status=("Done {0}/{1}\t".format(cls.__processorCount, cls.__processorTotalCount)) )

    @classmethod
    def endProcessorCountBar(cls):
        cls.__processorCount = 0
        cls.__bShowProcessBar = False
        cmds.progressWindow(endProgress=1)

    @classmethod
    def checkIfUseProcessorCountBar(cls,inTotalCount):
        # type: (int) -> bool
        cls.endProcessorCountBar()
        cls.__bShowProcessBar = True if inTotalCount > cls.vtxUpdateThresholdNum else False
        if cls.__bShowProcessBar:
            cls.setProcessorCountBarTotalNum(inTotalCount)
        return cls.__bShowProcessBar

    @classmethod
    def getCurProcessorCountBarStatus (cls):
        # type: (...) -> bool
        return cls.__bShowProcessBar

    @classmethod
    def getCurProcessorCountBarNum(cls):
        # type: (...) -> int
        return cls.__processorCount

class MMeshFaceVtxProcessorInterface():
    __debug = False
    halfVec = api.MVector(0.5, 0.5, 0.5)

    @classmethod
    def iterMeshFaceVtx (cls, inDagPath, inColorSetName = "FaceVtxNorCol"):
        # type: (api.MDagPath, str) -> bool
        return False

class MMeshFaceVtxProcessorBase(MMeshFaceVtxProcessorInterface):
    
    __metaclass__ = MSingleton
    __processorBarCls = MMProcessorBar
    
    @classmethod
    def iterMeshFaceVtx (cls, inDagPath, inColorSetName = "FaceVtxNorCol"):
        # type: (api.MDagPath, str) -> bool
        fnMesh = api.MFnMesh(inDagPath)

        # NOTE: Init processor bar
        cls.__processorBarCls.checkIfUseProcessorCountBar(fnMesh.numFaceVertices)

        # NOTE: force set current color set
        if not cls._processCurColorSet(fnMesh, inColorSetName): return False
        
        # NOTE: Traverse all faceVertex in current mesh
        meshFaceVertexIt = api.MItMeshFaceVertex(inDagPath)

        return cls._iterMeshFaceVtxDataProcess(fnMesh, meshFaceVertexIt)
    
    @classmethod
    def _processCurColorSet(cls, fnMesh, inColorSetName):
        # type: (api.MFnMesh, str) -> bool
        """ check current color set """
        return False

    @classmethod
    def _iterMeshFaceVtxDataProcess (cls, fnMesh, meshFaceVertexIt):
        #type: (api.MFnMesh, api.MItMeshFaceVertex,) -> bool
        """Traverse all faceVertex in current mesh"""
        return False

class MMeshFaceVtxProcessorCol2Nor(MMeshFaceVtxProcessorBase):
    
    @classmethod
    def _processCurColorSet(cls, fnMesh, inColorSetName):
        # type: (api.MFnMesh, str) -> bool
        """ check current color set """
        
        # NOTE: check if valid color set to set color
        colorSets = fnMesh.getColorSetNames() # type: list[str]
        if not inColorSetName in colorSets:
            # NOTE: Create new color set
            api.MGlobal.displayWarning(u"\"{}\" dont have colorSet \"{}\" to recovery normal".format(fnMesh.dagPath(), inColorSetName))
            return False
        return True

    @classmethod
    def _iterMeshFaceVtxDataProcess (cls, fnMesh, meshFaceVertexIt):
        #type: (api.MFnMesh, api.MItMeshFaceVertex,) -> bool
        """Traverse all faceVertex in current mesh"""

        # NOTE: Get Current Color Set name
        curColorSetName = fnMesh.currentColorSetName()
        
        # NOTE: Save to temp list
        normals = api.MVectorArray()
        faceIds = api.MIntArray()
        vertexIds=api.MIntArray()

        # NOTE: Processor bar update data
        curIndex = 0
        nextUpdateNum = cls.__processorBarCls.vtxUpdateThresholdNum
        
        while not meshFaceVertexIt.isDone():
            faceVtxCol = meshFaceVertexIt.getColor(curColorSetName) # type: api.MColor
            faceVtxCor2Vec = api.MVector(faceVtxCol[0],faceVtxCol[1],faceVtxCol[2])
            normals.append((faceVtxCor2Vec - cls.halfVec) * 2)
            vertexIds.append(meshFaceVertexIt.vertexId())
            faceIds.append(meshFaceVertexIt.faceId())            
            meshFaceVertexIt.next()
            
            # NOTE: update processor bar
            curIndex+=1
            if curIndex >= nextUpdateNum:
                nextUpdateNum += cls.__processorBarCls.vtxUpdateThresholdNum
                cls.__processorBarCls.updateProcessorCountBar(curIndex)
        else:
            fnMesh.setFaceVertexNormals(normals, faceIds, vertexIds)
            normals = faceIds = vertexIds = None
            cls.__processorBarCls.endProcessorCountBar() # NOTE: end processor bar
            return True
        return False
    
class MMeshFaceVtxProcessorWsNor2Col(MMeshFaceVtxProcessorBase):
        
    @classmethod
    def _processCurColorSet(cls,fnMesh, inColorSetName):
        # type: (api.MFnMesh, str) -> bool
        """ force set current color set """
        
        # NOTE: check if valid color set to set color
        colorSets = fnMesh.getColorSetNames() # type: list[str]
        if not inColorSetName in colorSets:
            # NOTE: Create new color set
            fnMesh.createColorSet(inColorSetName, True)
        fnMesh.setCurrentColorSetName(inColorSetName)
        return True

    @classmethod
    def _iterMeshFaceVtxDataProcess (cls, fnMesh, meshFaceVertexIt):
        #type: (api.MFnMesh, api.MItMeshFaceVertex,) -> bool
        """Traverse all faceVertex in current mesh"""
        
        # -- NOTE: Save to temp list
        colors = api.MColorArray()
        faceIds = api.MIntArray()
        vertexIds=api.MIntArray()
        
        # NOTE: Processor bar update data
        curIndex = 0
        nextUpdateNum = cls.__processorBarCls.vtxUpdateThresholdNum
        modifier = api.MDGModifier()
        
        while not meshFaceVertexIt.isDone():
            norVec = meshFaceVertexIt.getNormal() # type: api.MVector
            colors.append(norVec * 0.5 + cls.halfVec)
            vertexIds.append(meshFaceVertexIt.vertexId())
            faceIds.append(meshFaceVertexIt.faceId())
            meshFaceVertexIt.next()
            
            # NOTE: update processor bar
            curIndex+=1
            if curIndex >= nextUpdateNum:
                nextUpdateNum += cls.__processorBarCls.vtxUpdateThresholdNum
                cls.__processorBarCls.updateProcessorCountBar(curIndex)
        else:
            fnMesh.setFaceVertexColors(colors, faceIds, vertexIds, modifier=modifier)
            modifier.doIt()
            colors = faceIds = vertexIds = None
            cls.__processorBarCls.endProcessorCountBar() # NOTE: end processor bar
            return True
        return False

class MMeshFaceVtxProcessorTanNor2Col(MMeshFaceVtxProcessorInterface):
    
    @classmethod
    def iterMeshFaceVtx (cls, inDagPath, inColorSetName = "FaceVtxNorCol"):
        # type: (api.MDagPath, str) -> bool
        modifier = api.MDGModifier()
        fnMesh = api.MFnMesh(inDagPath)
        
        # NOTE: check if valid color set to set color
        colorSets = fnMesh.getColorSetNames() # type: list[str]
        if not inColorSetName in colorSets:
            # NOTE: Create new color set
            fnMesh.createColorSet(inColorSetName, True)
        fnMesh.setCurrentColorSetName(inColorSetName)

        # NOTE: Get current uvSet name
        curUvSet = fnMesh.currentUVSetName() # type: str
        
        # NOTE: Traverse all faceVertex in current mesh
        meshFaceVertexIt = api.MItMeshFaceVertex(inDagPath)

        # -- NOTE: Save to temp list
        colors = api.MColorArray()
        faceIds = api.MIntArray()
        vertexIds=api.MIntArray()
        zeroArray = (0,0,0)
        
        while not meshFaceVertexIt.isDone():
            norVec = meshFaceVertexIt.getNormal() # type: api.MVector
            binorVec = meshFaceVertexIt.getBinormal(uvSet=curUvSet)  # type: api.MVector
            tanVec = meshFaceVertexIt.getTangent(uvSet=curUvSet)  # type: api.MVector
            
            # NOTE: build TBN Matrix
            tbnMatrix = api.MMatrix([vec[i] if i < 3 else 0 for vec in (tanVec, binorVec, norVec, zeroArray) for i in range(4)])
            
            # NOTE: pre-multiply TBN matrix. same as post-multiply transposed TBN matrix
            tanNorVec =  tbnMatrix * norVec  # type: api.MVector
            colors.append(tanNorVec * 0.5 + cls.halfVec)
            vertexIds.append(meshFaceVertexIt.vertexId())
            faceIds.append(meshFaceVertexIt.faceId())
            meshFaceVertexIt.next()
        else:
            fnMesh.setFaceVertexColors(colors, faceIds, vertexIds, modifier=modifier)
            modifier.doIt()
            colors = faceIds = vertexIds = None
            return True
        return False

class MVtxNorColConverter (object):
    __debug = True
    __convertClass = MMeshFaceVtxProcessorInterface
    __validConvertClass = False
    __colorSetName = "FaceVtxNorCol"

    @classmethod
    def setConvertClass (cls, inConvertClass):
        # type: (MMeshFaceVtxProcessorInterface) -> bool
        if inConvertClass.__name__ in (subCls.__name__ for subCls in MMeshFaceVtxProcessorInterface.__subclasses__()):
            cls.__convertClass = inConvertClass
            cls.__validConvertClass = True
            return True
        else:
            raise Exception("Not valid convert class")

    @classmethod
    def setColorSetName (cls, colSetName):
        # type: (str) -> bool
        if isinstance(colSetName,str):
            cls.__colorSetName = colSetName
            return True
        return False
    
    @classmethod
    def getColorSetName (cls):
        # type: (...) -> str
        return cls.__colorSetName

    @classmethod
    def getCurConvertClassName (cls):
        # type: (...) -> str
        return cls.__convertClass.__name__
    
    @classmethod
    def convertSelMesh (cls):
        # type: (...) -> bool
        curMeshList = MSelectionHelper.getCurMeshSelections()
        if curMeshList.isEmpty(): 
            api.MGlobal.displayWarning(u"No mesh has been selected?")
            return False
        
        curMeshListIt = api.MItSelectionList(MSelectionHelper.getCurMeshSelections())
        while not (curMeshListIt.isDone()):
            curMeshDagPath= curMeshListIt.getDagPath() # type: api.MDagPath
            if not cls.iterMeshFaceVtx(curMeshDagPath): print(u"process mesh \"{}\" fail".format(curMeshDagPath))
            if cmds.progressWindow( query=True, isCancelled=True ):
                api.MGlobal.displayWarning("Processor end by user!")
                break
            curMeshListIt.next()
        else:
            return True

    @classmethod
    def iterMeshFaceVtx (cls, inDagPath):
        # type: (api.MDagPath) -> bool
        if not cls.__validConvertClass :
            raise Exception("Not set a valid convert class")
        return cls.__convertClass.iterMeshFaceVtx(inDagPath, cls.__colorSetName)

def runScript_nor2col(*args, **kw):
    MVtxNorColConverter.setConvertClass(MMeshFaceVtxProcessorWsNor2Col)
    MVtxNorColConverter.convertSelMesh()

def runScript_tanNor2col(*args, **kw):
    """Generate Tangent space Normal, test only!"""
    MVtxNorColConverter.setConvertClass(MMeshFaceVtxProcessorTanNor2Col)
    MVtxNorColConverter.convertSelMesh()

def runScript_col2nor(*args, **kw):
    MVtxNorColConverter.setConvertClass(MMeshFaceVtxProcessorCol2Nor)
    MVtxNorColConverter.convertSelMesh()

if __name__ == "__main__":
    runScript_nor2col()