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