From 8de2d826d3200d032bfc9a509e09ebf33f7c5064 Mon Sep 17 00:00:00 2001 From: Juksten <58948183+Juksten@users.noreply.github.com> Date: Mon, 16 Dec 2019 20:09:40 +0100 Subject: [PATCH] Update tensionmap.py --- tensionmap.py | 172 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 104 insertions(+), 68 deletions(-) diff --git a/tensionmap.py b/tensionmap.py index a5aee8c..b90e8d9 100644 --- a/tensionmap.py +++ b/tensionmap.py @@ -16,6 +16,8 @@ # along with this program. If not, see . import bpy +import numpy as np +from numba import jit,njit bl_info = { "name": "Tension Map Script", @@ -122,80 +124,117 @@ def tm_update(obj, context): depsgraph = context.evaluated_depsgraph_get() object_eval = obj.evaluated_get(depsgraph) deformed_mesh = object_eval.to_mesh() + mesh = obj.data # restore modifiers viewport show state for i in range_modifiers: obj.modifiers[i].show_viewport = show_original_state[i] - # array to store new weight for each vertices - weights = [0.0] * num_vertices - - # calculate the new weights - for i in range(len(obj.data.edges)): - edge = obj.data.edges[i] - first_vertex = edge.vertices[0] - second_vertex = edge.vertices[1] - - original_edge_length = (obj.data.vertices[first_vertex].co - - obj.data.vertices[second_vertex].co).length - deformed_edge_length = ( - deformed_mesh.vertices[first_vertex].co - deformed_mesh.vertices[second_vertex].co).length - - deformation_factor = (original_edge_length - - deformed_edge_length) * obj.data.tm_multiply - - # store the weights by subtracting to overlay all the factors for each vertex - weights[first_vertex] -= deformation_factor - weights[second_vertex] -= deformation_factor + edges = read_edges(mesh) + + verts_original = read_vertices(mesh) + + verts_deformed = read_vertices(deformed_mesh) + + original_edge_len = np.linalg.norm(verts_original[edges[:,0]] + - verts_original[edges[:,1]],axis=1) + + deformed_edge_len = np.linalg.norm(verts_deformed[edges[:,0]] + - verts_deformed[edges[:,1]],axis=1) + + deform_factor = (original_edge_len - deformed_edge_len) * obj.data.tm_multiply + + weights = np.zeros(num_vertices,dtype='float32') + + get_weights(weights,edges,deform_factor,num_vertices) # delete the temporary deformed mesh object_eval.to_mesh_clear() - # create vertex color list for faster access only if vertex color is activated - if obj.data.tm_enable_vertex_colors: - vertex_colors = [0.0] * (number_of_tm_channels * num_vertices) - - # calculate the new values - # store them in the vertex_colors array if the feature is active - # store them in the vertex groups if the feature is active - for i in range(num_vertices): - stretch_value = obj.data.tm_minimum - squeeze_value = obj.data.tm_minimum - - if weights[i] >= 0: - # positive: stretched - stretch_value = max(obj.data.tm_minimum, min( - obj.data.tm_maximum, weights[i])) - else: - # negative: squeezed - # invert weights to keep only positive values - squeeze_value = max(obj.data.tm_minimum, min( - obj.data.tm_maximum, -weights[i])) - - if obj.data.tm_enable_vertex_groups: - add_index = [i] - group_squeeze.add(add_index, squeeze_value, "REPLACE") - group_stretch.add(add_index, stretch_value, "REPLACE") - - if obj.data.tm_enable_vertex_colors: - # red - vertex_colors[i * number_of_tm_channels] = stretch_value - # green - vertex_colors[i * number_of_tm_channels + 1] = squeeze_value - - # store the calculated vertex colors if the feature is active + stretch = np.zeros(num_vertices,dtype='float32') + squeeze = np.zeros(num_vertices,dtype='float32') + + calc_values(stretch,squeeze,num_vertices,weights,0,1) + + if obj.data.tm_enable_vertex_groups: + for i in range(num_vertices): + group_squeeze.add([i], squeeze[i], "REPLACE") + group_stretch.add([i], stretch[i], "REPLACE") + if obj.data.tm_enable_vertex_colors: - colors_tension = get_or_create_vertex_colors(obj, "tm_tension") - # this is heavy, but vertex colors are stored by vertex loop - # and there is no simpler way to do it (it would seem) - for poly_idx in range(len(obj.data.polygons)): - polygon = obj.data.polygons[poly_idx] - for loop_vertex_idx, loop_idx in enumerate(polygon.loop_indices): - vertex_color = colors_tension.data[loop_idx] - vertex_idx = polygon.vertices[loop_vertex_idx] - # replace the color by a 4D vector, using 0 for blue and 1 for alpha - vertex_color.color = (vertex_colors[vertex_idx * number_of_tm_channels], - vertex_colors[vertex_idx * number_of_tm_channels + 1], 0, 1) + if 'tm_tension' not in obj.data.vertex_colors: + obj.data.vertex_colors.new(name='tm_tension') + + vertex_colors = np.zeros(num_vertices*number_of_tm_channels,dtype='float32') + + get_colors(vertex_colors,num_vertices,number_of_tm_channels,stretch,squeeze) + + vertex_colors_data = read_vertex_color_data(mesh) + + vertex_idxs = read_polygon_vertices(mesh) + + num_polygons = len(mesh.polygons) + + loop_indices = np.arange(num_polygons*4).reshape(num_polygons,4) + + change_colors(mesh,num_polygons,loop_indices,vertex_colors_data,vertex_idxs,vertex_colors,number_of_tm_channels) + + vertex_colors_data = vertex_colors_data.reshape(len(mesh.polygons)*16) + + mesh.vertex_colors['tm_tension'].data.foreach_set('color',vertex_colors_data) + +def read_vertex_color_data(mesh): + vertex_colors = np.zeros(len(mesh.polygons)*16, dtype='float32') + mesh.vertex_colors['tm_tension'].data.foreach_get('color',vertex_colors) + return (vertex_colors.reshape(len(mesh.polygons)*4,4)) + +def read_polygon_vertices(mesh): + verts = np.zeros(len(mesh.polygons)*4, dtype='int32') + mesh.polygons.foreach_get('vertices',verts) + return(verts.reshape(len(mesh.polygons),4)) + +def read_vertices(mesh): + verts = np.zeros((len(mesh.vertices)*3), dtype='float32') + mesh.vertices.foreach_get("co", verts) + return (verts.reshape(len(mesh.vertices), 3)) + +def read_edges(mesh): + edges = np.zeros((len(mesh.edges)*2), dtype='int32') + mesh.edges.foreach_get("vertices", edges) + return (edges.reshape(len(mesh.edges), 2)) + +@jit() +def change_colors(mesh,num_polygons,loop_indices,vertex_colors_data,vertex_idxs,vertex_colors,num_of_channels): + for poly_idx in range(num_polygons): + for loop_vertex_idx, loop_idx in enumerate(loop_indices[poly_idx,:]): + vertex_idx = vertex_idxs[poly_idx,loop_vertex_idx] + vertex_colors_data[loop_idx,:] = [vertex_colors[vertex_idx * 2], + vertex_colors[vertex_idx * 2 + 1], 0, 1] + +@njit() +def get_weights(wg,edges,deform_factor,num_vertices): + for i in range(len(deform_factor)): + for j in edges[i,:]: + wg[j] -= deform_factor[i] + +@njit() +def get_colors(vertex_colors,num_vertices,channels,stretch,squeeze): + for i in range(num_vertices): + vertex_colors[i * channels] = stretch[i] + vertex_colors[i * channels + 1] = squeeze[i] + +@njit() +def calc_values(stretch_values,squeeze_values,num_vertices,weights,min_val,max_val): + for i in range(num_vertices): + if weights[i] >= 0: + stretch_values[i] = max(min_val, min( + max_val, weights[i])) + squeeze_values[i] = min_val + else: + squeeze_values[i] = max(min_val, min( + max_val, -weights[i])) + stretch_values[i] = min_val + def tm_update_handler(scene): @@ -240,7 +279,7 @@ def execute(self, context): return {"FINISHED"} def invoke(self, context, event): - return self.execute(context) + return self.execute(context) class TmPanel(bpy.types.Panel): @@ -279,14 +318,11 @@ def draw(self, context): ''' # TODO: finish implementing interface for choosing modifiers flow.separator() - row2 = flow.column() row2.enabled = context.object.data.tm_active row2.label(text="Modifiers to use when computing tension") list = row2.box() - modifiers = context.object.modifiers - for i in range(len(modifiers)): row = list.row() row.prop(modifiers[i], "use_for_tension", text=modifiers[i].name)