LeitLightLightLeitLeidLite/simulate.py

#!/usr/bin/env python3

# Ugly, but working:
#
# 1. Start blender
# 2. Import your light strip with one vertex per LED, ensure vertex index matches LED index
# 3. Load this script in blender and run it
# 4. Start animation
# 5. Switch to material preview or rendered view
# 6. Broadcast your WLED DNRGB updates to the config below

#FIXME: This might need a bit of https://b3d.interplanety.org/en/accessing-custom-attributes-created-in-geometry-nodes/ to get the input indices

HOST = "localhost"
PORT = 21324

#import LeitLightLightLeitLeidLite as l6

import bpy
import mathutils

import socket



obj = bpy.data.objects['strips']

if True:
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((HOST, PORT))
    sock.setblocking(False)

ledCount = len(obj.data.vertices)

print("%d LEDs" % ledCount)
state = bytearray([0x00,0x00,0x00] * ledCount)
durations = bytearray([0x00] * ledCount)


mesh = obj.data #bpy.context.collection.objects["strips"].data
def setAttrib(attribName, data):
    if attribName in mesh.attributes:
        attribute = mesh.attributes[attribName]
    else:
        attribute = mesh.attributes.new(name=attribName, type="FLOAT", domain="POINT")
    #attribute_values = [i for i in range(len(mesh.vertices))]
    attribute.data.foreach_set("value", data)

# Created using https://extensions.blender.org/add-ons/node-to-python/
def createMaterial(materialName):

  mat = bpy.data.materials.new(name = materialName)
  mat.use_nodes = True

  material = mat.node_tree
  #start with a clean node tree
  for node in material.nodes:
      material.nodes.remove(node)
  material.color_tag = 'NONE'
  material.description = ""
  material.default_group_node_width = 140
  
  #node Material Output
  material_output = material.nodes.new("ShaderNodeOutputMaterial")
  material_output.name = "Material Output"
  material_output.is_active_output = True
  material_output.target = 'ALL'
  #Displacement
  material_output.inputs[2].default_value = (0.0, 0.0, 0.0)
  #Thickness
  material_output.inputs[3].default_value = 0.0

  #node Emission
  emission = material.nodes.new("ShaderNodeEmission")
  emission.name = "Emission"
  #Color
  emission.inputs[0].default_value = (1.0, 0.0, 0.018193421885371208, 1.0)

  #node Add Shader
  add_shader = material.nodes.new("ShaderNodeAddShader")
  add_shader.name = "Add Shader"

  #node Emission.001
  emission_001 = material.nodes.new("ShaderNodeEmission")
  emission_001.name = "Emission.001"
  #Color
  emission_001.inputs[0].default_value = (0.021969836205244064, 1.0, 0.0, 1.0)

  #node Add Shader.001
  add_shader_001 = material.nodes.new("ShaderNodeAddShader")
  add_shader_001.name = "Add Shader.001"

  #node Emission.002
  emission_002 = material.nodes.new("ShaderNodeEmission")
  emission_002.name = "Emission.002"
  #Color
  emission_002.inputs[0].default_value = (0.0, 0.0016676115337759256, 1.0, 1.0)

  #node Attribute
  attribute = material.nodes.new("ShaderNodeAttribute")
  attribute.name = "Attribute"
  attribute.attribute_name = "red"
  attribute.attribute_type = 'GEOMETRY'

  #node Attribute.001
  attribute_001 = material.nodes.new("ShaderNodeAttribute")
  attribute_001.name = "Attribute.001"
  attribute_001.attribute_name = "green"
  attribute_001.attribute_type = 'GEOMETRY'

  #node Attribute.002
  attribute_002 = material.nodes.new("ShaderNodeAttribute")
  attribute_002.name = "Attribute.002"
  attribute_002.attribute_name = "blue"
  attribute_002.attribute_type = 'GEOMETRY'


  #Set locations
  material_output.location = (300.0, 300.0)
  emission.location = (-90.9997329711914, 320.24908447265625)
  add_shader.location = (119.9999771118164, 290.62481689453125)
  emission_001.location = (-95.54425048828125, 215.4246063232422)
  add_shader_001.location = (115.45545959472656, 163.01229858398438)
  emission_002.location = (-100.0887680053711, 87.81208801269531)
  attribute.location = (-358.0585021972656, 399.44940185546875)
  attribute_001.location = (-350.34014892578125, 246.54885864257812)
  attribute_002.location = (-354.8846740722656, 114.37872314453125)

  #Set dimensions
  material_output.width, material_output.height = 140.0, 100.0
  emission.width, emission.height = 140.0, 100.0
  add_shader.width, add_shader.height = 140.0, 100.0
  emission_001.width, emission_001.height = 140.0, 100.0
  add_shader_001.width, add_shader_001.height = 140.0, 100.0
  emission_002.width, emission_002.height = 140.0, 100.0
  attribute.width, attribute.height = 140.0, 100.0
  attribute_001.width, attribute_001.height = 140.0, 100.0
  attribute_002.width, attribute_002.height = 140.0, 100.0

  #emission.Emission -> add_shader.Shader
  material.links.new(emission.outputs[0], add_shader.inputs[0])
  #emission_001.Emission -> add_shader.Shader
  material.links.new(emission_001.outputs[0], add_shader.inputs[1])
  #emission_002.Emission -> add_shader_001.Shader
  material.links.new(emission_002.outputs[0], add_shader_001.inputs[1])
  #add_shader.Shader -> add_shader_001.Shader
  material.links.new(add_shader.outputs[0], add_shader_001.inputs[0])
  #attribute.Fac -> emission.Strength
  material.links.new(attribute.outputs[2], emission.inputs[1])
  #attribute_002.Fac -> emission_002.Strength
  material.links.new(attribute_002.outputs[2], emission_002.inputs[1])
  #attribute_001.Fac -> emission_001.Strength
  material.links.new(attribute_001.outputs[2], emission_001.inputs[1])
  #add_shader_001.Shader -> material_output.Surface
  material.links.new(add_shader_001.outputs[0], material_output.inputs[0])
  return mat


# Created using https://extensions.blender.org/add-ons/node-to-python/
def createGeoNodes(geoNodeName, material):
  geometry_nodes = bpy.data.node_groups.new(type = 'GeometryNodeTree', name = geoNodeName)

  geometry_nodes.color_tag = 'NONE'
  geometry_nodes.description = ""
  geometry_nodes.default_group_node_width = 140
  

  geometry_nodes.is_modifier = True

  #geometry_nodes interface
  #Socket Geometry
  geometry_socket = geometry_nodes.interface.new_socket(name = "Geometry", in_out='OUTPUT', socket_type = 'NodeSocketGeometry')
  geometry_socket.attribute_domain = 'POINT'

  #Socket Geometry
  geometry_socket_1 = geometry_nodes.interface.new_socket(name = "Geometry", in_out='INPUT', socket_type = 'NodeSocketGeometry')
  geometry_socket_1.attribute_domain = 'POINT'


  #initialize geometry_nodes nodes
  #node Group Input
  group_input = geometry_nodes.nodes.new("NodeGroupInput")
  group_input.name = "Group Input"

  #node Group Output
  group_output = geometry_nodes.nodes.new("NodeGroupOutput")
  group_output.name = "Group Output"
  group_output.is_active_output = True

  #node Instance on Points
  instance_on_points = geometry_nodes.nodes.new("GeometryNodeInstanceOnPoints")
  instance_on_points.name = "Instance on Points"
  #Selection
  instance_on_points.inputs[1].default_value = True
  #Pick Instance
  instance_on_points.inputs[3].default_value = False
  #Instance Index
  instance_on_points.inputs[4].default_value = 0
  #Rotation
  instance_on_points.inputs[5].default_value = (0.0, 0.0, 0.0)
  #Scale
  instance_on_points.inputs[6].default_value = (1.0, 1.0, 1.0)

  #node Ico Sphere
  ico_sphere = geometry_nodes.nodes.new("GeometryNodeMeshIcoSphere")
  ico_sphere.name = "Ico Sphere"
  #Radius
  ico_sphere.inputs[0].default_value = 1.0
  #Subdivisions
  ico_sphere.inputs[1].default_value = 1

  #node Store Named Attribute
  store_named_attribute = geometry_nodes.nodes.new("GeometryNodeStoreNamedAttribute")
  store_named_attribute.name = "Store Named Attribute"
  store_named_attribute.data_type = 'INT'
  store_named_attribute.domain = 'POINT'
  #Selection
  store_named_attribute.inputs[1].default_value = True
  #Name
  store_named_attribute.inputs[2].default_value = "index"

  #node Index
  index = geometry_nodes.nodes.new("GeometryNodeInputIndex")
  index.name = "Index"

  #node Realize Instances
  realize_instances = geometry_nodes.nodes.new("GeometryNodeRealizeInstances")
  realize_instances.name = "Realize Instances"
  #Selection
  realize_instances.inputs[1].default_value = True
  #Realize All
  realize_instances.inputs[2].default_value = True
  #Depth
  realize_instances.inputs[3].default_value = 0

  #node Capture Attribute
  capture_attribute = geometry_nodes.nodes.new("GeometryNodeCaptureAttribute")
  capture_attribute.name = "Capture Attribute"
  capture_attribute.active_index = 0
  capture_attribute.capture_items.clear()
  capture_attribute.capture_items.new('FLOAT', "Index")
  capture_attribute.capture_items["Index"].data_type = 'INT'
  capture_attribute.domain = 'POINT'

  #node Set Material
  set_material = geometry_nodes.nodes.new("GeometryNodeSetMaterial")
  set_material.name = "Set Material"
  #Selection
  set_material.inputs[1].default_value = True
  set_material.inputs[2].default_value = material





  #Set locations
  group_input.location = (-537.77783203125, 192.87576293945312)
  group_output.location = (1020.145263671875, 338.443115234375)
  instance_on_points.location = (-40.82560348510742, 80.0461654663086)
  ico_sphere.location = (-380.4179382324219, -85.53570556640625)
  store_named_attribute.location = (503.62579345703125, 345.2210388183594)
  index.location = (-583.5319213867188, 81.38147735595703)
  realize_instances.location = (235.0821075439453, 147.38453674316406)
  capture_attribute.location = (-276.59991455078125, 312.9664306640625)
  set_material.location = (815.2620239257812, 334.15179443359375)

  #Set dimensions
  group_input.width, group_input.height = 140.0, 100.0
  group_output.width, group_output.height = 140.0, 100.0
  instance_on_points.width, instance_on_points.height = 140.0, 100.0
  ico_sphere.width, ico_sphere.height = 140.0, 100.0
  store_named_attribute.width, store_named_attribute.height = 140.0, 100.0
  index.width, index.height = 140.0, 100.0
  realize_instances.width, realize_instances.height = 140.0, 100.0
  capture_attribute.width, capture_attribute.height = 140.0, 100.0
  set_material.width, set_material.height = 140.0, 100.0

  #initialize geometry_nodes links
  #ico_sphere.Mesh -> instance_on_points.Instance
  geometry_nodes.links.new(ico_sphere.outputs[0], instance_on_points.inputs[2])
  #group_input.Geometry -> capture_attribute.Geometry
  geometry_nodes.links.new(group_input.outputs[0], capture_attribute.inputs[0])
  #capture_attribute.Geometry -> instance_on_points.Points
  geometry_nodes.links.new(capture_attribute.outputs[0], instance_on_points.inputs[0])
  #index.Index -> capture_attribute.Index
  geometry_nodes.links.new(index.outputs[0], capture_attribute.inputs[1])
  #capture_attribute.Index -> store_named_attribute.Value
  geometry_nodes.links.new(capture_attribute.outputs[1], store_named_attribute.inputs[3])
  #instance_on_points.Instances -> realize_instances.Geometry
  geometry_nodes.links.new(instance_on_points.outputs[0], realize_instances.inputs[0])
  #realize_instances.Geometry -> store_named_attribute.Geometry
  geometry_nodes.links.new(realize_instances.outputs[0], store_named_attribute.inputs[0])
  #store_named_attribute.Geometry -> set_material.Geometry
  geometry_nodes.links.new(store_named_attribute.outputs[0], set_material.inputs[0])
  #set_material.Geometry -> group_output.Geometry
  geometry_nodes.links.new(set_material.outputs[0], group_output.inputs[0])
  return geometry_nodes





def my_handler(scene):
  global state
  print("Frame Change", scene.frame_current)


  try:
    data, addr = sock.recvfrom(4096)
    #print("received message: %s" % data)

    assert(data[0] == 4) # DNRGB
    duration = data[1]
    baseIndex = (data[2] << 8) | data[3]
    payloadLength = len(data) - 4
    assert(payloadLength % 3 == 0)
    count = payloadLength // 3
    newState = data[4:]
    

    beforeCount = baseIndex
    afterIndex = baseIndex + count
    state = state[0:beforeCount*3] + newState + state[afterIndex*3:]

  except BlockingIOError:
    pass


  setAttrib('red', [x / 255.0 for x in state[0::3]])
  setAttrib('green', [x / 255.0 for x in state[1::3]])
  setAttrib('blue', [x / 255.0 for x in state[2::3]])

def regHandler(my_handler):
  for func in bpy.app.handlers.frame_change_post[:]:
      if func.__name__ == my_handler.__name__:
          bpy.app.handlers.frame_change_post.remove(func)
  bpy.app.handlers.frame_change_pre.append(my_handler)

material = createMaterial("l6-geometry-nodes")
geometry_nodes = createGeoNodes("l6-material", material)

mod = obj.modifiers.new('l6-modifier', 'NODES')
mod.node_group = geometry_nodes

regHandler(my_handler)