set_foliage_vertex_data.py

import bpy

FOLIAGE_DATA_VERTEX_COLOR_MAP = "FoliageData"
        
def get_wind_influence(obj, vert):
    map_index = obj.data.vertex_colors.find(FOLIAGE_DATA_VERTEX_COLOR_MAP)
    map = obj.data.vertex_colors[map_index]

    for polygon in obj.data.polygons:
        for i, index in enumerate(polygon.vertices):
            if vert.index == index:
                loop_index = polygon.loop_indices[i]
                vertex_data = map.data[loop_index]                        
                return vertex_data.color[0] # The wind influence is stored in the red channel, ie. channel 0.

def set_wind_influence(obj, verts, wind_influence):
    map_index = obj.data.vertex_colors.find(FOLIAGE_DATA_VERTEX_COLOR_MAP)
    map = obj.data.vertex_colors[map_index]
        
    for polygon in obj.data.polygons:
        for vert in verts:
            for i, index in enumerate(polygon.vertices):
                if vert.index == index:
                    loop_index = polygon.loop_indices[i]
                    vertex_data = map.data[loop_index]
                    vertex_data.color[0] = wind_influence
    

class SetFoliageVertexData(bpy.types.Operator):
    """Set Foliage Vertex Data"""
    bl_idname = "basis_tools.set_foliage_vertex_data"
    bl_label = "Set Foliage Vertex Data"
    
    wind_influence_prop = bpy.props.FloatProperty(name = "Wind Influence", min = 0.0, max = 1.0)
    
    prev_object_mode = ""
    selected_verts = []

    @classmethod
    def poll(cls, context):
        return context.active_object is not None

    def execute(self, context):
        set_wind_influence(context.active_object, self.selected_verts, self.wind_influence_prop)
        bpy.ops.object.mode_set(mode = self.prev_object_mode)
        return {'FINISHED'}
    
    def invoke(self, context, event):
        if len(bpy.context.selected_objects) != 1:
            self.report({'WARNING'}, "Cannot set foliage vertex data. Select exactly one object.")
            return {'CANCELLED'}
        
        obj = bpy.context.active_object
        
        if obj.type != "MESH":
            self.report({'WARNING'}, "Cannot set foliage vertex data. The selected object is not a mesh.")
            return {'CANCELLED'}
        
        if obj.data.vertex_colors.find(FOLIAGE_DATA_VERTEX_COLOR_MAP) == -1:
            self.report({'WARNING'}, \
                "Cannot set foliage vertex data. The selected object does not have a vertex color map named '" + \
                FOLIAGE_DATA_VERTEX_COLOR_MAP + "'.")
            return {'CANCELLED'}
        
        # We need to switch from Edit mode to Object mode so the selection gets updated.
        self.prev_object_mode = obj.mode
        bpy.ops.object.mode_set(mode='OBJECT')
        
        self.selected_verts = [v for v in obj.data.vertices if v.select]
        
        if len(self.selected_verts) == 0:
            self.report({'WARNING'}, "Cannot set foliage vertex data. No vertices selected.")
            return {'CANCELLED'}
        
        self.wind_influence_prop = get_wind_influence(obj, self.selected_verts[0])
        
        wm = context.window_manager
        return wm.invoke_props_dialog(self)

def register():
    bpy.utils.register_class(SetFoliageVertexData)

def unregister():
    bpy.utils.unregister_class(SetFoliageVertexData)