0%

Maya中法线转顶点色且可逆

發表於 2023-02-26 分類於水

摘要

因为项目需要, 编写了一个Maya脚本, 可以快速将顶点(切线/世界)法线转换为顶点色, 并可逆转换回来.

基本基于OpenMaya API 1.0编写, 速度非常快.只需要几秒就可以处理100W+三角形

实现代码

人生苦短
我用Python
除非性能不够

# coding=utf-8

# Script by MineClever, help to process high mesh normal on low mesh ~
# Bug report to : https://t.me/minemineSlaver

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

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(six.with_metaclass(MSingleton,object)):
    __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(six.with_metaclass(MSingleton,object)):
    __debug = False
    
    __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(six.with_metaclass(MSingleton,object)):
    __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):
    
    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 MMeshFaceVtxProcessorTanOffsetNor2Col(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)
        refColorSetName= inColorSetName
        
        while not meshFaceVertexIt.isDone():
            faceVtxCol = meshFaceVertexIt.getColor(refColorSetName) # type: api.MColor
            faceVtxCor2Vec = api.MVector(faceVtxCol[0],faceVtxCol[1],faceVtxCol[2])
            decodedNorVec = (faceVtxCor2Vec - cls.halfVec) * 2 # type: api.MVector
            
            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 * decodedNorVec  # 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 MMeshFaceVtxProcessorInterface.__subclasscheck__(inConvertClass):
            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(curMeshList)
        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 convertSelMesh (cls):
        try:
            cls._convertSelMesh()
        except Exception as E:
            api.MGlobal.displayError(str(E))
        finally:
            # NOTE: Force end process window
            cmds.progressWindow(endProgress=1)

    @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_tanNor2tanCol(*args, **kw):
    """Generate Tangent space Normal, test only!"""
    MVtxNorColConverter.setConvertClass(MMeshFaceVtxProcessorTanOffsetNor2Col)
    MVtxNorColConverter.convertSelMesh()

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

if __name__ == "__main__":
    runScript_nor2col()

注释

理论上还是够快的
如果用C++ 重写, 应该可以考虑上多线程, 分割处理所有顶点面, 应当可以处理上亿三角形

歡迎關注我的其它發布渠道