Baking AO and curvature to vertex colour

BakeToVC.py

import bpy
import sys
import os

aoVC = None
aovcName = "AOVC"
pointinessMatName = "NewMat"
cVC = None
cvcName = "CVC"
finalVCName = "FinalVC"
fbxFilePath = ""
existingVC = None

def getFileSelectionFromArg():
    global fbxFilePath
    if sys.argv[4]:
        fbxFilePath = sys.argv[4]

def objectSelection():
    # TODO: Change this based on needs. Get all meshes in hierarchy
    obj = bpy.context.active_object
    return obj

def clearVC():
    # Delete all vertex color except if its already existed before baking
    obj = objectSelection()
    existingVC = GetExistingVC()
    for col_attr in obj.data.vertex_colors:
        try:
            if col_attr.name not in existingVC.name:
                obj.data.vertex_colors.remove(col_attr)
                obj.update_tag()
        except:
            continue
    for attr in obj.data.color_attributes:
        try:
            if attr.name not in existingVC.name:
                obj.data.color_attributes.remove(attr)
                obj.update_tag()
        except:
            continue
    for area in bpy.context.screen.areas:
        area.tag_redraw()

def AOBake():
    obj = objectSelection()
    # Add vertex color attributes in Color Attributes
    mesh = obj.data
    #if not mesh.vertex_colors:
    aovc = mesh.vertex_colors.new()
    aovc.name = aovcName

    # Set active color attribute to AO
    namedCA = bpy.context.object.data.color_attributes
    setACA = namedCA.get(aovc.name)
    bpy.context.object.data.attributes.active_color = setACA

    bpy.context.scene.render.engine = 'CYCLES'
    bpy.context.scene.render.bake.target = 'VERTEX_COLORS'
    bpy.context.scene.cycles.device = 'GPU'
    # TODO: Add extra settings here
    # Do render
    bpy.ops.object.bake(type='AO')

def cacheMaterial(obj):
    materialCache = []
    materialCache.clear()
    for materialSlot in obj.material_slots:
        if materialSlot.material is not None:
            materialCache.append(materialSlot.material)
    return materialCache

def deleteAllOldPointinessMaterial():
    allMat = bpy.data.materials
    for m in allMat:
        if pointinessMatName in m:
            bpy.data.materials.remove(m)

def createPointinessMaterial():    
    # Clean up unused material
    deleteAllOldPointinessMaterial()
    #bpy.ops.outliner.orphans_purge()
    obj = objectSelection()

    for matSlot in obj.material_slots:
        try:
            mName = matSlot.name
            if (pointinessMatName in mName):
                bpy.ops.object.material_slot_remove({'object': obj})
                bpy.data.materials.remove(matSlot.material)
        except:
            continue

    # Then make NewMat
    newMat = bpy.data.materials.new(name=pointinessMatName)
    newMat.use_nodes = True
    materialNodes = newMat.node_tree.nodes
    materialLinks = newMat.node_tree.links
    for node in materialNodes:
        materialNodes.remove(node)
    # Create a Geometry node and add it to the material
    geometryNode = materialNodes.new(type="ShaderNodeNewGeometry")
    geometryNode.location = (0, 0)
    # Add output node
    outputNode = materialNodes.new(type='ShaderNodeOutputMaterial')
    outputNode.location = (600, 0)
    # Add color Ramp
    colorRamp = materialNodes.new(type="ShaderNodeValToRGB")
    colorRamp.color_ramp.elements[0].position = 0.45
    colorRamp.color_ramp.elements[1].position = 0.6
    colorRamp.location = (200, 0)
    # Connect them all
    materialLinks.new(geometryNode.outputs[7], colorRamp.inputs[0])
    materialLinks.new(colorRamp.outputs[0], outputNode.inputs[0])

def createBetterPointinessMaterial():
    # Clean up unused material
    deleteAllOldPointinessMaterial()
    #bpy.ops.outliner.orphans_purge()
    obj = objectSelection()

    for matSlot in obj.material_slots:
        try:
            mName = matSlot.name
            if (pointinessMatName in mName):
                bpy.ops.object.material_slot_remove({'object': obj})
                bpy.data.materials.remove(matSlot.material)
        except:
            continue

    # Then make NewMat
    newMat = bpy.data.materials.new(name=pointinessMatName)
    newMat.use_nodes = True
    materialNodes = newMat.node_tree.nodes
    materialLinks = newMat.node_tree.links
    for node in materialNodes:
        materialNodes.remove(node)

    # Value node
    edgeValue = materialNodes.new(type="ShaderNodeValue")        
    edgeValue.location = (-100,0)
    edgeValue.outputs[0].default_value = 0.025
    geometryNode = materialNodes.new(type="ShaderNodeNewGeometry")
    geometryNode.location = (0, 0)
    aoNode = materialNodes.new(type="ShaderNodeAmbientOcclusion")
    aoNode.location = (200, 700)
    aoNode.samples = 32
    aoNode.only_local = True
    # Add color Ramp
    colorRamp = materialNodes.new(type="ShaderNodeValToRGB")
    colorRamp.location = 400, 600
    colorRamp.color_ramp.elements[0].position = 0.4

    # Add bevel node
    bevelNode = materialNodes.new(type="ShaderNodeBevel")
    bevelNode.location = (400, 0)
    # Add shader node vector math (dot product)
    dotProduct = materialNodes.new(type="ShaderNodeVectorMath")
    dotProduct.location = (600,200)
    # Invert color node
    invertNode = materialNodes.new(type="ShaderNodeInvert")
    invertNode.location = (800, -100)
    # 2nd Color Ramp
    secondColorRamp = materialNodes.new(type="ShaderNodeValToRGB")
    secondColorRamp.color_ramp.elements[0].position = 0
    secondColorRamp.color_ramp.elements[1].position = 0.4
    secondColorRamp.location = (1000, -500)
    secondColorRamp.color_ramp.elements[0].color = (0.3, 0.3, 0.3, 1)
    # Shader node mix / multiply
    mixNode = materialNodes.new(type="ShaderNodeMix")
    mixNode.location = (1200, -800)

    # Add output node
    outputNode = materialNodes.new(type='ShaderNodeOutputMaterial')
    outputNode.location = (1600, 0)

    dotProduct.operation = 'DOT_PRODUCT'
    mixNode.blend_type = 'MULTIPLY'
    mixNode.data_type = 'RGBA'
    mixNode.inputs[0].default_value = 1.0

    # CONNECT THESE NODES
    materialLinks.new(edgeValue.outputs[0], aoNode.inputs[1])
    materialLinks.new(geometryNode.outputs[3], aoNode.inputs[2])
    materialLinks.new(edgeValue.outputs[0], bevelNode.inputs[0])
    materialLinks.new(aoNode.outputs[1], colorRamp.inputs[0])
    materialLinks.new(bevelNode.outputs[0], dotProduct.inputs[0])
    materialLinks.new(geometryNode.outputs[3], dotProduct.inputs[1])
    materialLinks.new(dotProduct.outputs[1], invertNode.inputs[1])
    materialLinks.new(invertNode.outputs[0], secondColorRamp.inputs[0])
    materialLinks.new(colorRamp.outputs[0], mixNode.inputs[6])
    materialLinks.new(secondColorRamp.outputs[0], mixNode.inputs[-1])
    materialLinks.new(mixNode.outputs[-1], outputNode.inputs[0])

def assignPointinessMaterial(obj):
    if len(obj.material_slots) == 0:
        bpy.ops.object.material_slot_add()
    for slot in obj.material_slots:
        if pointinessMatName in bpy.data.materials:
            material = bpy.data.materials[pointinessMatName]
            slot.material = material

def bakeCurvatureToVC():
    print ("Curvature bake")
    obj = objectSelection()
    # Add vertex color attributes in Color Attributes
    mesh = obj.data
    cVC = mesh.vertex_colors.new()
    cVC.name = cvcName

    # Set active color attribute to AO
    namedCA = bpy.context.object.data.color_attributes
    setACA = namedCA.get(cvcName)
    bpy.context.object.data.attributes.active_color = setACA

    bpy.context.scene.render.engine = 'CYCLES'
    bpy.context.scene.render.bake.target = 'VERTEX_COLORS'
    bpy.context.scene.cycles.device = 'GPU'

    bpy.ops.object.bake(type='EMIT')

def reapplyMaterial(mc):
    obj = objectSelection()
    if obj.data.materials:
        for index in range(len(obj.material_slots)):
            obj.material_slots[index].material = mc[index]

def CurvatureBake():
    obj = objectSelection()
    # Cache existing material
    materialCache = cacheMaterial(obj)
    # Create new material with pointiness Geometry node
    createPointinessMaterial()
    # Assign pointiness material to all material slots in an object
    assignPointinessMaterial(obj)
    bakeCurvatureToVC()
    reapplyMaterial(materialCache)

def CopyAttributesToVC():
    obj = objectSelection()
    mesh = obj.data
    fVC = mesh.vertex_colors.new()
    fVC.name = finalVCName

    # Get source and target vertex color
    ao = obj.data.vertex_colors.get(aovcName)
    curvature = obj.data.vertex_colors.get(cvcName)
    target = obj.data.vertex_colors.get(finalVCName)
    try:
        existing = obj.data.vertex_colors.get(existingVC.name)
        # Copy existing vertex color to Red Channel
        for loopIndex, targetLoop in enumerate(target.data):
            sourceColor = existing.data[loopIndex].color
            targetColor = targetLoop.color
            targetColor[0] = sourceColor[0]
            try:
                targetColor[3] = sourceColor[3]
            except:
                continue
        obj.data.update()
    except:
        print("No existing VC")

    # Copy attribute Curvature to Blue channel
    for loopIndex, targetLoop in enumerate(target.data):
        sourceColor = curvature.data[loopIndex].color
        targetColor = targetLoop.color
        targetColor[2] = sourceColor[0]
    obj.data.update()

    # Copy attribute AO to Green channel
    for loopIndex, targetLoop in enumerate(target.data):
        sourceColor = ao.data[loopIndex].color
        targetColor = targetLoop.color
        targetColor[1] = sourceColor[0]
    obj.data.update()

def DeleteUnusedAttributes():
    # Remove extra attributes
    obj = objectSelection()
    for vcAttr in obj.data.color_attributes:
        try:
            if (vcAttr.name == cvcName) :
                obj.data.color_attributes.remove(vcAttr)
            if (vcAttr.name == aovcName) :
                obj.data.color_attributes.remove(vcAttr)
            if (vcAttr.name == existingVC.name) :
                obj.data.color_attributes.remove(vcAttr)
        except:
            continue
    for vcAttr in obj.data.vertex_colors:
        try:
            if (vcAttr.name == cvcName) :
                obj.data.vertex_colors.remove(vcAttr)
            if (vcAttr.name == aovcName) :
                obj.data.vertex_colors.remove(vcAttr)
            if (vcAttr.name == existingVC.name) :
                obj.data.color_attributes.remove(vcAttr)
        except:
            continue

def PrepareScene():    
    try:
        if (sys.argv[4]):
            # Select all objects in the scene and delete
            bpy.ops.object.select_all(action='SELECT')
            bpy.ops.object.delete()
            ImportFBX()
    except:
        print("No fbx to open")

def ImportFBX():
    getFileSelectionFromArg()
    bpy.ops.import_scene.fbx(filepath=fbxFilePath)    
    imported_objects = bpy.context.selected_objects
    for obj in imported_objects:
        obj.select_set(True)
    
    bpy.context.view_layer.objects.active = imported_objects[-1]

def ExportFBXSettings():
    bpy.ops.export_scene.fbx(
        filepath=fbxFilePath,
        use_selection=True,
        axis_forward='Z',
        axis_up='Y',
        bake_anim=False,
        bake_anim_use_all_actions=False
    )

def ExportFBX():
    # TODO: Check export settings
    print ("Export FBX")
    obj = objectSelection()
    obj.select_set(True)
    if os.path.exists(fbxFilePath) and os.access(fbxFilePath, os.W_OK):
        ExportFBXSettings()

    else:
        # If file is read only, change the permission
        os.chmod(fbxFilePath, 0o644)

        ExportFBXSettings()

def QuitBlender():
    bpy.ops.object.select_all(action='DESELECT')
    bpy.ops.wm.quit_blender()

def GetExistingVC():
    global existingVC
    obj = objectSelection()
    try:
        vc = obj.data.color_attributes[0]
    except:
        vc = None
    if vc and (existingVC == None):
        existingVC = vc
        return vc

def main():    
    PrepareScene()
    clearVC()
    AOBake()
    CurvatureBake()
    CopyAttributesToVC()
    DeleteUnusedAttributes()
    ExportFBX()
    QuitBlender()

if __name__ == "__main__":
    main()