5e2fff7e41d4b4a65730dc4f5a447e234451f158
[blender-addons-contrib.git] / mesh_show_vgroup_weights.py
1 # ***** BEGIN GPL LICENSE BLOCK *****
2 #
3 #
4 # This program is free software; you can redistribute it and/or
5 # modify it under the terms of the GNU General Public License
6 # as published by the Free Software Foundation; either version 2
7 # of the License, or (at your option) any later version.
8 #
9 # This program is distributed in the hope that it will be useful,
10 # but WITHOUT ANY WARRANTY; without even the implied warranty of
11 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12 # GNU General Public License for more details.
13 #
14 # You should have received a copy of the GNU General Public License
15 # along with this program; if not, write to the Free Software Foundation,
16 # Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17 #
18 # ***** END GPL LICENCE BLOCK *****
19
20 # <pep8 compliant> (Thanks to CodemanX on IRC)
21
22 bl_info = {
23     "name": "Show Vertex Groups/Weights",
24     "author": "Jason van Gumster (Fweeb), Bartius Crouch",
25     "version": (0, 7, 1),
26     "blender": (2, 62, 3),
27     "location": "3D View > Properties Region > Show Weights",
28     "description": "Finds the vertex groups of a selected vertex and displays the corresponding weights",
29     "warning": "Requires bmesh",
30     "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/Py/"\
31         "Scripts/Modeling/Show_Vertex_Group_Weights",
32     "tracker_url": "http://projects.blender.org/tracker/index.php?"\
33         "func=detail&aid=30609",
34     "category": "Mesh"}
35
36 #TODO - Add button for selecting vertices with no groups
37
38
39 import bpy, bmesh, bgl, blf, mathutils
40
41
42 # Borrowed/Modified from Bart Crouch's old Index Visualizer add-on
43 def calc_callback(self, context):
44     #polling
45     if context.mode != "EDIT_MESH" or len(context.active_object.vertex_groups) == 0:
46         return
47
48     # get color info from theme
49     acol = context.user_preferences.themes[0].view_3d.editmesh_active
50     tcol = (acol[0] * 0.85, acol[1] * 0.85, acol[2] * 0.85)
51     
52     # get screen information
53     mid_x = context.region.width / 2.0
54     mid_y = context.region.height / 2.0
55     width = context.region.width
56     height = context.region.height
57     
58     # get matrices
59     view_mat = context.space_data.region_3d.perspective_matrix
60     ob_mat = context.active_object.matrix_world
61     total_mat = view_mat * ob_mat
62     
63     # calculate location info
64     texts = []
65     locs = []
66     weights  = []
67     me = context.active_object.data
68     bm = bmesh.from_edit_mesh(me)
69     dvert_lay = bm.verts.layers.deform.active
70
71     for v in bm.verts:
72         if v.select: #XXX Should check v.hide here, but it doesn't work
73             if bm.select_mode == {'VERT'} and bm.select_history.active is not None and bm.select_history.active.index == v.index:
74                 locs.append([acol[0], acol[1], acol[2], v.index, v.co.to_4d()])
75             else:
76                 locs.append([tcol[0], tcol[1], tcol[2], v.index, v.co.to_4d()])
77             dvert = v[dvert_lay]
78             for vgroup in context.active_object.vertex_groups:
79                 if vgroup.index in dvert.keys():
80                     weights += [v.index, vgroup.index, dvert[vgroup.index]]
81
82     for loc in locs:
83         vec = total_mat * loc[4] # order is important
84         # dehomogenise
85         vec = mathutils.Vector((vec[0] / vec[3], vec[1] / vec[3], vec[2] / vec[3]))
86         x = int(mid_x + vec[0] * width / 2.0)
87         y = int(mid_y + vec[1] * height / 2.0)
88         texts += [loc[0], loc[1], loc[2], loc[3], x, y, 0]
89
90     # store as ID property in mesh
91     context.active_object.data["show_vgroup_verts"] = texts
92     context.active_object.data["show_vgroup_weights"] = weights
93
94
95 # draw in 3d-view
96 def draw_callback(self, context):
97     # polling
98     if context.mode != "EDIT_MESH" or len(context.active_object.vertex_groups) == 0:
99         return
100     # retrieving ID property data
101     try:
102         texts = context.active_object.data["show_vgroup_verts"]
103         weights = context.active_object.data["show_vgroup_weights"]
104     except:
105         return
106     if not texts:
107         return
108
109     bm = bmesh.from_edit_mesh(context.active_object.data)
110
111     if bm.select_mode == {'VERT'} and bm.select_history.active is not None:
112         active_vert = bm.select_history.active
113     else:
114         active_vert = None
115
116     # draw
117     blf.size(0, 13, 72)
118     blf.enable(0, blf.SHADOW)
119     blf.shadow(0, 3, 0.0, 0.0, 0.0, 1.0)
120     blf.shadow_offset(0, 2, -2)
121     for i in range(0, len(texts), 7):
122         bgl.glColor3f(texts[i], texts[i+1], texts[i+2])
123         blf.position(0, texts[i+4], texts[i+5], texts[i+6])
124         blf.draw(0, "Vertex " + str(int(texts[i+3])) + ":")
125         font_y = texts[i+5]
126         group_name = ""
127         for j in range(0, len(weights), 3):
128             if int(weights[j]) == int(texts[i+3]):
129                 font_y -= 13
130                 blf.position(0, texts[i+4] + 10, font_y, texts[i+6])
131                 for group in context.active_object.vertex_groups:
132                     if group.index == int(weights[j+1]):
133                         group_name = group.name
134                         break
135                 blf.draw(0, group_name + ": %.3f" % weights[j+2])
136         if group_name == "":
137             font_y -= 13
138             blf.position(0, texts[i+4] + 10, font_y, texts[i+6])
139             blf.draw(0, "No Groups")
140
141     # restore defaults
142     blf.disable(0, blf.SHADOW)
143
144
145 # operator
146 class ShowVGroupWeights(bpy.types.Operator):
147     bl_idname = "view3d.show_vgroup_weights"
148     bl_label = "Show Vertex Group Weights"
149     bl_description = "Toggle the display of the vertex groups and weights for selected vertices"
150     
151     @classmethod
152     def poll(cls, context):
153         return context.mode == 'EDIT_MESH'
154     
155     def __del__(self):
156         bpy.context.scene.display_indices = -1
157         clear_properties(full=False)
158     
159     def modal(self, context, event):
160         if context.area:
161             context.area.tag_redraw()
162
163         # removal of callbacks when operator is called again
164         if context.scene.display_indices == -1:
165             context.region.callback_remove(self.handle1)
166             context.region.callback_remove(self.handle2)
167             context.scene.display_indices = 0
168             return {'CANCELLED'}
169         
170         return {'PASS_THROUGH'}
171     
172     def invoke(self, context, event):
173         if context.area.type == 'VIEW_3D':
174             if context.scene.display_indices < 1:
175                 # operator is called for the first time, start everything
176                 context.scene.display_indices = 1
177                 self.handle1 = context.region.callback_add(calc_callback,
178                     (self, context), 'POST_VIEW')
179                 self.handle2 = context.region.callback_add(draw_callback,
180                     (self, context), 'POST_PIXEL')
181                 context.window_manager.modal_handler_add(self)
182                 return {'RUNNING_MODAL'}
183             else:
184                 # operator is called again, stop displaying
185                 context.scene.display_indices = -1
186                 clear_properties(full=False)
187                 return {'RUNNING_MODAL'}
188         else:
189             self.report({'WARNING'}, "View3D not found, can't run operator")
190             return {'CANCELLED'}
191
192
193 # properties used by the script
194 class InitProperties(bpy.types.Operator):
195     bl_idname = "view3d.init_find_weights"
196     bl_label = "Initialize properties for vgroup weights finder"
197     
198     def execute(self, context):
199         bpy.types.Scene.display_indices = bpy.props.IntProperty(
200             name="Display indices",
201             default=0)
202         context.scene.display_indices = 0
203         return {'FINISHED'}
204
205
206 # removal of ID-properties when script is disabled
207 def clear_properties(full=True):
208     # can happen on reload
209     if bpy.context.scene is None:
210         return
211     
212     if "show_vgroup_verts" in bpy.context.active_object.data.keys():
213         del bpy.context.active_object.data["show_vgroup_verts"]
214     if "show_vgroup_weights" in bpy.context.active_object.data.keys():
215         del bpy.context.active_object.data["show_vgroup_weights"]
216     if full:
217         props = ["display_indices"]
218         for p in props:
219             if p in bpy.types.Scene.bl_rna.properties:
220                 exec("del bpy.types.Scene." + p)
221             if p in bpy.context.scene.keys():
222                 del bpy.context.scene[p]
223
224
225 class AssignVertexWeight(bpy.types.Operator):
226     bl_idname = "mesh.vertex_group_assign"
227     bl_label = "Assign Weights"
228     bl_description = "Assign weights for all of the groups on a specific vertex"
229
230     vgroup_weights = bpy.props.StringProperty(name = "Vertex Group Weights")
231
232     @classmethod
233     def poll(cls, context):
234         return context.mode == 'EDIT_MESH'
235
236     def execute(self, context):
237         me = context.active_object.data
238         bm = bmesh.from_edit_mesh(me)
239         dvert_lay = bm.verts.layers.deform.active
240         weights = eval(self.vgroup_weights) #XXX Would be nice if I didn't have to use an eval
241
242         for v in bm.verts:
243             if v.index == weights["__index__"]:
244                 del weights["__index__"]
245                 dvert = v[dvert_lay]
246                 for vgroup in dvert.keys():
247                     dvert[vgroup] = weights[vgroup]
248                 break
249
250         return {'FINISHED'}
251
252
253 class RemoveFromVertexGroup(bpy.types.Operator):
254     bl_idname = "mesh.vertex_group_remove"
255     bl_label = "Remove Vertex from Group"
256     bl_description = "Remove a specific vertex from a specific vertex group"
257
258     #XXX abusing vector props here a bit; the first element is the vert index and the second is the group index
259     vert_and_group = bpy.props.IntVectorProperty(name = "Vertex and Group to remove", size = 2)
260
261     @classmethod
262     def poll(cls, context):
263         return context.mode == 'EDIT_MESH'
264
265     def execute(self, context):
266         ob = context.active_object
267         me = ob.data
268         bm = bmesh.from_edit_mesh(me)
269
270         # Save current selection
271         selected_verts = []
272         for v in bm.verts:
273             if v.select is True:
274                 selected_verts.append(v.index)
275                 if v.index != self.vert_and_group[0]:
276                     v.select = False
277
278         ob.vertex_groups.active_index = self.vert_and_group[1]
279         bpy.ops.object.vertex_group_remove_from()
280
281         # Re-select vertices
282         for v in bm.verts:
283             if v.index in selected_verts:
284                 v.select = True
285
286         #XXX Hacky, but there's no other way to update the UI panels
287         bpy.ops.object.editmode_toggle()
288         bpy.ops.object.editmode_toggle()
289         return {'FINISHED'}
290
291
292 class AddToVertexGroup(bpy.types.Operator):
293     bl_idname = "mesh.vertex_group_add"
294     bl_label = "Add Vertex to Group"
295     bl_description = "Add a specific vertex to a specific vertex group"
296
297     def avail_vgroups(self, context):
298         ob = context.active_object
299         bm = bmesh.from_edit_mesh(ob.data)
300         dvert_lay = bm.verts.layers.deform.active
301         items = []
302         self.vertex = bm.select_history.active.index
303
304         dvert = bm.select_history.active[dvert_lay]
305
306         items.append(("-1", "New Vertex Group", "-1", -1))
307
308         for i in ob.vertex_groups:
309             if i.index not in dvert.keys():
310                 items.append((i.name, i.name, str(i.index), i.index))
311
312         return items
313
314     vertex = bpy.props.IntProperty()
315     available_vgroups = bpy.props.EnumProperty(items = avail_vgroups, name = "Available Groups")
316
317     @classmethod
318     def poll(cls, context):
319         return context.mode == 'EDIT_MESH'
320
321     def execute(self, context):
322         ob = context.active_object
323         me = ob.data
324         bm = bmesh.from_edit_mesh(me)
325         #print(self.available_vgroups)
326
327         # Save current selection
328         selected_verts = []
329         for v in bm.verts:
330             if v.select is True:
331                 selected_verts.append(v.index)
332                 if v.index != self.vertex:
333                     v.select = False
334
335         weight = context.tool_settings.vertex_group_weight
336         context.tool_settings.vertex_group_weight = 1.0
337         if self.available_vgroups == "-1":
338             bpy.ops.object.vertex_group_assign(new = True) #XXX Assumes self.vertex is the active vertex
339         else:
340             bpy.ops.object.vertex_group_set_active(group = self.available_vgroups)
341             bpy.ops.object.vertex_group_assign() #XXX Assumes self.vertex is the active vertex
342         context.tool_settings.vertex_group_weight = weight
343
344         # Re-select vertices
345         for v in bm.verts:
346             if v.index in selected_verts:
347                 v.select = True
348
349         #XXX Hacky, but there's no other way to update the UI panels
350         bpy.ops.object.editmode_toggle()
351         bpy.ops.object.editmode_toggle()
352         return {'FINISHED'}
353
354
355 class PanelShowWeights(bpy.types.Panel):
356     bl_label = "Show Weights"
357     bl_space_type = "VIEW_3D"
358     bl_region_type = "UI"
359     bl_options = {'DEFAULT_CLOSED'}
360
361     @classmethod
362     def poll(cls, context):
363         return context.mode == 'EDIT_MESH'
364
365     def draw(self, context):
366         layout = self.layout
367         ob = context.active_object
368         me = ob.data
369         bm = bmesh.from_edit_mesh(me)
370         dvert_lay = bm.verts.layers.deform.active
371
372         if context.scene.display_indices < 1:
373             layout.operator(ShowVGroupWeights.bl_idname, text = "Show Weights Overlay")
374         else:
375             layout.operator(ShowVGroupWeights.bl_idname, text = "Hide Weights Overlay")
376
377         if len(ob.vertex_groups) > 0:
378             # Active vertex
379             active_vert = bm.select_history.active
380             sub = layout.box()
381             col = sub.column(align = True)
382             if bm.select_mode == {'VERT'} and active_vert is not None:
383                 col.label(text = "Active Vertex")
384                 row = col.row()
385                 row.label(text = "Vertex " + str(active_vert.index) + ":")
386                 row.operator_menu_enum("mesh.vertex_group_add", "available_vgroups", text = "Add Group", icon = 'GROUP_VERTEX')
387                 has_groups = False
388                 vgroup_weights = {}
389
390                 for i in me.vertices:
391                     if i.index == active_vert.index:
392                         vgroup_weights["__index__"] = i.index
393                         for j in range(len(i.groups)):
394                             for k in ob.vertex_groups:
395                                 if k.index == i.groups[j].group:
396                                     has_groups = True
397                                     split = col.split(percentage = 0.90, align = True)
398                                     vgroup_weights[k.index] = i.groups[j].weight
399                                     row = split.row(align = True)
400                                     row.prop(i.groups[j], "weight", text = k.name, slider = True, emboss = not k.lock_weight)
401                                     row = split.row(align = True)
402                                     row.operator("mesh.vertex_group_remove", text = "R").vert_and_group = (i.index, k.index)
403                 
404                 if not has_groups:
405                     col.label(text = "    No Groups")
406                 else:
407                     col.operator("mesh.vertex_group_assign").vgroup_weights = str(vgroup_weights)
408                 layout.separator()
409             else:
410                 col.label(text = "No Active Vertex")
411             layout.prop(context.window_manager, "show_vgroups_show_all", toggle = True)
412             # All selected vertices (except for the active vertex)
413             if context.window_manager.show_vgroups_show_all:
414                 for v in bm.verts:
415                     if v.select:
416                         if active_vert is not None and v.index == active_vert.index:
417                             continue
418                         sub = layout.box()
419                         col = sub.column(align = True)
420                         col.label(text = "Vertex " + str(v.index) + ":")
421                         has_groups = False
422                         vgroup_weights = {}
423                         for i in me.vertices:
424                             if i.index == v.index:
425                                 vgroup_weights["__index__"] = i.index
426                                 for j in range(len(i.groups)):
427                                     for k in ob.vertex_groups:
428                                         if k.index == i.groups[j].group:
429                                             has_groups = True
430                                             split = col.split(percentage = 0.90, align = True)
431                                             vgroup_weights[k.index] = i.groups[j].weight
432                                             row = split.row(align = True)
433                                             row.prop(i.groups[j], "weight", text = k.name, slider = True, emboss = not k.lock_weight)
434                                             row = split.row(align = True)
435                                             row.operator("mesh.vertex_group_remove", text = "R").vert_and_group = (i.index, k.index)
436                         if not has_groups:
437                             col.label(text = "    No Groups")
438                         else:
439                             col.operator("mesh.vertex_group_assign").vgroup_weights = str(vgroup_weights)
440         else:
441             layout.label(text = "No Groups")
442
443
444 def register():
445     bpy.types.WindowManager.show_vgroups_show_all = bpy.props.BoolProperty(
446         name = "Show All Selected Vertices",
447         description = "Show all vertices with vertex groups assigned to them",
448         default = False)
449     bpy.types.Mesh.assign_vgroup = bpy.props.StringProperty()
450     bpy.utils.register_class(ShowVGroupWeights)
451     bpy.utils.register_class(InitProperties)
452     bpy.ops.view3d.init_find_weights()
453     bpy.utils.register_class(AssignVertexWeight)
454     bpy.utils.register_class(RemoveFromVertexGroup)
455     bpy.utils.register_class(AddToVertexGroup)
456     bpy.utils.register_class(PanelShowWeights)
457     
458
459 def unregister():
460     bpy.utils.unregister_class(ShowVGroupWeights)
461     bpy.utils.unregister_class(InitProperties)
462     clear_properties()
463     bpy.utils.unregister_class(AssignVertexWeight)
464     bpy.utils.unregister_class(RemoveFromVertexGroup)
465     bpy.utils.unregister_class(AddToVertexGroup)
466     bpy.utils.unregister_class(PanelShowWeights)
467     del bpy.types.WindowManager.show_vgroups_show_all
468     del bpy.types.Mesh.assign_vgroup
469
470 if __name__ == "__main__":
471     register()