Nuke - Lightsmith Normalisation

lightsmishNormalisation.py

def getWeightedAverage(node):
    width, height = node.width(), node.height()
    red = green = blue = 0.0
    for x in range(width):
        for y in range(height):
            red += node.sample("red", x, y)
            green += node.sample("green", x, y)
            blue += node.sample("blue", x, y)
    count = width * height
    return red / count, green / count, blue / count

def setLightsmithNormalisation(node):
    saturation = nuke.nodes.Saturation(saturation=0)
    saturation.setInput(0, node)
    
    average = getWeightedAverage(saturation)
    normalisation = 1 / ((average[0] + average[1] + average[2]) / 3)
    
    multiply = nuke.nodes.Multiply(name="Lightsmith_Multiply", value=normalisation)
    multiply.setInput(0, node)

    write = nuke.nodes.Write(name="Lightsmith_Write")
    write.setInput(0, multiply)
    
    nuke.delete(saturation)
    for node in (write, multiply):
        nuke.autoplace(multiply)

    average =  getWeightedAverage(multiply)
    normalisation = max(average)
    print("*" * 80)
    print("Lightsmith Parameters:")
    print("LScolor = %s" % ",".join(map(lambda x: str(int((x / normalisation) * 255)), average)))
    print("LSmulti = %s" % normalisation)
    print("*" * 80)

    return True

setLightsmithNormalisation(nuke.selectedNode())