Ambient Occlusion to vertex color channel

ambient-occlusion-vertex-bake.py

import math
import sys
import bpy
import numpy as np
import traceback
import mathutils
from mathutils.bvhtree import BVHTree

def print(data):
    for window in bpy.context.window_manager.windows:
        screen = window.screen
        for area in screen.areas:
            if area.type == 'CONSOLE':
                override = {'window': window, 'screen': screen, 'area': area}
                bpy.ops.console.scrollback_append(override, text=str(data), type="OUTPUT")

ob = bpy.context.active_object

print('---')

verbose = True

def describeThing(thing):
    pr('describing thing:')
    pr('  str:'+str(thing))
    pr('  type: ' + str(type(thing)))

def gammaCorrect(color, gamma):
    if(not isinstance(color, list)):
        pr('??'+str(color))
#        color = list(color)
#    color = coerceColor(color)
    color[0] = pow(color[0], gamma)
    color[1] = pow(color[1], gamma)
    color[2] = pow(color[2], gamma)
    return color


def coerceColor(val):
    if(isinstance(val, float)):
        return [val, val, val, 1.0]
    else:
        c = list(val)
        if(len(c) == 3):
            c.append(1.0)
        str('-------'+str(c))
        return c

def coerceFac(val):
    fac = 0.0
    if(not isinstance(val, float)):
#        describeThing(val)
        for i in range(3):
            fac += val[i]
        fac /= 3.0
    else:
        fac = val
    return fac

def lerp(vals1, vals2, amt):
    if(isinstance(vals1, np.float64)):
        vals1 = [vals1, vals1, vals1, 1.0]
    if(isinstance(vals2, np.float64)):
        vals2 = [vals2, vals2, vals2, 1.0]
    while(len(vals1) < len(vals2)):
        vals1 = np.append(vals1, 1.0)
    while(len(vals1) > len(vals2)):
        vals2 = np.append(vals2, 1.0)
        
    amt = coerceFac(amt)
        
    result = []
    for i in range(3):
        result.append(vals1[i] * (1.0-amt) + vals2[i] * amt)
    return result

def pr(msg):
    if(verbose):
        print(msg)

def err(msg):
    pr('ERROR: ' + msg)
    raise Exception(msg)


gr=(math.sqrt(5.0) + 1.0) / 2.0;  # golden ratio = 1.6180339887498948482
ga=(2.0 - gr) * 6.2831853072;  # golden angle = 2.39996322972865332

def fibonacci_spiral_sphere(i, num_points):
    lat = math.asin(-1.0 + 2.0 * i / (num_points+1.0));
    lon = ga * i;

    x = math.cos(lon)*math.cos(lat);
    y = math.sin(lon)*math.cos(lat);
    z = math.sin(lat);
    return [x, y, z]

def mix(a, b, amt):
    return np.add(np.multiply(b, amt), np.multiply(a, (1.0 - amt)))


tree = BVHTree.FromObject(ob, bpy.context.evaluated_depsgraph_get())
verts = ob.data.vertices
def shadow(attrLoop):
    loop = ob.data.loops[attrLoop]
    dir = list(loop.normal)
    dirEul = mathutils.Euler(dir, 'XYZ')
    averageNormal = list(verts[loop.vertex_index].normal)
    pos = list(verts[loop.vertex_index].co)
    offsetNormal = mix(averageNormal, dir, 0.0)
    offset = np.multiply(offsetNormal, 0.001)
    pos = np.add(pos, offset)
    shade = 0.0
    for i in range(100):
        fuzzyEul = mathutils.Euler(fibonacci_spiral_sphere(i, 200), 'XYZ')
        fuzzyQuat = fuzzyEul.to_quaternion()
        dirQuat = dirEul.to_quaternion()
        dirQuat.rotate(fuzzyQuat)
        finalDir = dirQuat.to_euler()
        hit, normal, index, distance = tree.ray_cast(pos, finalDir, 0.25)
        if(not hit == None):
            shade += 1.0
        
    light = 1.0 - shade / 100.0
#    return dir
    return [light, light, light, 1,0]


verbose = False

#print('vertex color channel: ' + exporterMatNode.name)


def dataTransferColor(loops, dstName, singleChannel = None):
    dst = ob.data.vertex_colors[dstName]
    
    for l in range(loops):
        cSrc = shadow(l)
#        cSrc = gammaCorrect(cSrc, 0.4545)
        cDst = dst.data[l].color
        describeThing(cSrc)
        if(singleChannel == None):
            cDst[0] = cSrc[0]
            cDst[1] = cSrc[1]
            cDst[2] = cSrc[2]
        else:
            fac = coerceFac(cSrc)
            cDst[singleChannel] = fac
            


#vertex colors are in fact stored per loop-vertex -&gt; MeshLoopColorLayer
if ob.type == 'MESH':
    
    ob.data.calc_normals_split()
    
    #how many loops do we have ?
    loops = len(ob.data.loops)
    print(str(loops))
   
    #go through each vertex color layer
    for vcol in ob.data.vertex_colors:
        print('vertex color channel: ' + vcol.name)
    
    
    dataTransferColor(loops, 'shadows')