mesh_selectbuffer: fix category, panel and context: T70017
[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, CoDEmanX",
25     "version": (0, 7, 2),
26     "blender": (2, 80, 4),
27     "location": "3D View > Properties Region > Show Weights",
28     "description": "Finds the vertex groups of a selected vertex "
29         "and displays the corresponding weights",
30     "warning": "",
31     "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/Py/"
32         "Scripts/Modeling/Show_Vertex_Group_Weights",
33     "tracker_url": "https://developer.blender.org/maniphest/task/edit/form/2/",
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.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     verts_max = context.scene.show_vgroups_weights_limit
72     verts_used = []
73
74     for elem in reversed(bm.select_history):
75         if not isinstance(elem, bmesh.types.BMVert): #or elem.hide:
76             continue
77         if bm.select_history.active == elem:
78             locs.append([acol[0], acol[1], acol[2], elem.index, elem.co.to_4d()])
79         else:
80             locs.append([tcol[0], tcol[1], tcol[2], elem.index, elem.co.to_4d()])
81         dvert = elem[dvert_lay]
82         for vgroup in context.active_object.vertex_groups:
83             if vgroup.index in dvert.keys():
84                 weights += [elem.index, vgroup.index, dvert[vgroup.index]]
85         verts_used.append(elem)
86         verts_max -= 1
87         if verts_max <= 0:
88             break
89
90     for v in bm.verts:
91         if v.select and v not in verts_used: #XXX Should check v.hide here, but it doesn't work
92             if isinstance(bm.select_history.active, bmesh.types.BMVert) and bm.select_history.active.index == v.index:
93                 locs.append([acol[0], acol[1], acol[2], v.index, v.co.to_4d()])
94             else:
95                 locs.append([tcol[0], tcol[1], tcol[2], v.index, v.co.to_4d()])
96             dvert = v[dvert_lay]
97             for vgroup in context.active_object.vertex_groups:
98                 if vgroup.index in dvert.keys():
99                     weights += [v.index, vgroup.index, dvert[vgroup.index]]
100             verts_max -= 1
101             if verts_max <= 0:
102                 break
103
104
105     for loc in locs:
106         vec = total_mat @ loc[4] # order is important
107         # dehomogenise
108         vec = mathutils.Vector((vec[0] / vec[3], vec[1] / vec[3], vec[2] / vec[3]))
109         x = int(mid_x + vec[0] * width / 2.0)
110         y = int(mid_y + vec[1] * height / 2.0)
111         texts += [loc[0], loc[1], loc[2], loc[3], x, y, 0]
112
113     # store as ID property in mesh
114     context.active_object.data["show_vgroup_verts"] = texts
115     context.active_object.data["show_vgroup_weights"] = weights
116
117
118 # draw in 3d-view
119 def draw_callback(self, context):
120     # polling
121     if context.mode != "EDIT_MESH" or len(context.active_object.vertex_groups) == 0:
122         return
123     # retrieving ID property data
124     try:
125         texts = context.active_object.data["show_vgroup_verts"]
126         weights = context.active_object.data["show_vgroup_weights"]
127     except:
128         return
129     if not texts:
130         return
131
132     bm = bmesh.from_edit_mesh(context.active_object.data)
133
134     if bm.select_mode == {'VERT'} and bm.select_history.active is not None:
135         active_vert = bm.select_history.active
136     else:
137         active_vert = None
138
139     # draw
140     blf.size(0, 13, 72)
141     blf.enable(0, blf.SHADOW)
142     blf.shadow(0, 3, 0.0, 0.0, 0.0, 1.0)
143     blf.shadow_offset(0, 2, -2)
144     for i in range(0, len(texts), 7):
145         bgl.glColor3f(texts[i], texts[i+1], texts[i+2])
146         blf.position(0, texts[i+4], texts[i+5], texts[i+6])
147         blf.draw(0, "Vertex " + str(int(texts[i+3])) + ":")
148         font_y = texts[i+5]
149         group_name = ""
150         for j in range(0, len(weights), 3):
151             if int(weights[j]) == int(texts[i+3]):
152                 font_y -= 13
153                 blf.position(0, texts[i+4] + 10, font_y, texts[i+6])
154                 for group in context.active_object.vertex_groups:
155                     if group.index == int(weights[j+1]):
156                         group_name = group.name
157                         break
158                 blf.draw(0, group_name + ": %.3f" % weights[j+2])
159         if group_name == "":
160             font_y -= 13
161             blf.position(0, texts[i+4] + 10, font_y, texts[i+6])
162             blf.draw(0, "No Groups")
163
164     # restore defaults
165     blf.disable(0, blf.SHADOW)
166
167
168 # operator
169 class ShowVGroupWeights(bpy.types.Operator):
170     bl_idname = "view3d.show_vgroup_weights"
171     bl_label = "Show Vertex Group Weights"
172     bl_description = "Toggle the display of the vertex groups and weights for selected vertices"
173
174     _handle_calc = None
175     _handle_draw = None
176
177     @staticmethod
178     def handle_add(self, context):
179         ShowVGroupWeights._handle_calc = bpy.types.SpaceView3D.draw_handler_add(
180             calc_callback, (self, context), 'WINDOW', 'POST_VIEW')
181         ShowVGroupWeights._handle_draw = bpy.types.SpaceView3D.draw_handler_add(
182             draw_callback, (self, context), 'WINDOW', 'POST_PIXEL')
183
184     @staticmethod
185     def handle_remove():
186         if ShowVGroupWeights._handle_calc is not None:
187             bpy.types.SpaceView3D.draw_handler_remove(ShowVGroupWeights._handle_calc, 'WINDOW')
188         if ShowVGroupWeights._handle_draw is not None:
189             bpy.types.SpaceView3D.draw_handler_remove(ShowVGroupWeights._handle_draw, 'WINDOW')
190         ShowVGroupWeights._handle_calc = None
191         ShowVGroupWeights._handle_draw = None
192
193     @classmethod
194     def poll(cls, context):
195         return context.mode == 'EDIT_MESH'
196
197     def execute(self, context):
198         if context.area.type == 'VIEW_3D':
199             if not context.scene.show_vgroups_weights:
200                 # operator is called and not running, start everything
201                 ShowVGroupWeights.handle_add(self, context)
202                 context.scene.show_vgroups_weights = True
203             else:
204                 # operator is called again, stop displaying
205                 ShowVGroupWeights.handle_remove()
206                 context.scene.show_vgroups_weights = False
207                 clear_properties(full=False)
208             context.area.tag_redraw()
209             return {'FINISHED'}
210         else:
211             self.report({'WARNING'}, "View3D not found, can't run operator")
212             return {'CANCELLED'}
213
214 class VGroupsWeights(bpy.types.PropertyGroup):
215     vgroup: bpy.props.IntProperty()
216     weight: bpy.props.FloatProperty(min=0.0, max=1.0)
217
218 class AssignVertexWeight(bpy.types.Operator):
219     bl_idname = "mesh.vertex_group_assign"
220     bl_label = "Assign Weights"
221     bl_description = "Assign weights for all of the groups on a specific vertex"
222
223     index: bpy.props.IntProperty()
224
225     vgroup_weights: bpy.props.CollectionProperty(
226         description="Vertex Group Weights",
227         type=VGroupsWeights)
228
229     @classmethod
230     def poll(cls, context):
231         return context.mode == 'EDIT_MESH'
232
233     def execute(self, context):
234         me = context.active_object.data
235         bm = bmesh.from_edit_mesh(me)
236         dvert_lay = bm.verts.layers.deform.active
237         weights = {}
238         for item in self.vgroup_weights:
239             weights[item.vgroup] = item.weight
240
241         for v in bm.verts:
242             if v.index == self.index:
243                 dvert = v[dvert_lay]
244                 for vgroup in dvert.keys():
245                     dvert[vgroup] = weights[vgroup]
246                 break
247         context.area.tag_redraw()
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 is 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         if context is None:
297             return []
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         #XXX since we need an identifier here, user won't be able to add a vgroup with that name ('-1')
307         #XXX could check against vgroup names and find an unused name, but it's a rare case after all.
308         items.append(("-1", "New Vertex Group", "Add a new vertex group to the active object", -1))
309
310         for i in ob.vertex_groups:
311             if i.index not in dvert.keys():
312                 items.append((i.name, i.name, str(i.index), i.index))
313
314         return items
315
316     vertex: bpy.props.IntProperty()
317     available_vgroups: bpy.props.EnumProperty(items=avail_vgroups, name="Available Groups")
318
319     @classmethod
320     def poll(cls, context):
321         return context.mode == 'EDIT_MESH'
322
323     def execute(self, context):
324         ob = context.active_object
325         me = ob.data
326         bm = bmesh.from_edit_mesh(me)
327         #print(self.available_vgroups)
328
329         # Save current selection
330         selected_verts = []
331         for v in bm.verts:
332             if v.select is True:
333                 selected_verts.append(v.index)
334                 if v.index != self.vertex:
335                     v.select = False
336
337         weight = context.tool_settings.vertex_group_weight
338         context.tool_settings.vertex_group_weight = 1.0
339         if self.available_vgroups == "-1":
340             bpy.ops.object.vertex_group_assign(new=True) #XXX Assumes self.vertex is the active vertex
341         else:
342             bpy.ops.object.vertex_group_set_active(group = self.available_vgroups)
343             bpy.ops.object.vertex_group_assign() #XXX Assumes self.vertex is the active vertex
344         context.tool_settings.vertex_group_weight = weight
345
346         # Re-select vertices
347         for v in bm.verts:
348             if v.index in selected_verts:
349                 v.select = True
350
351         #XXX Hacky, but there's no other way to update the UI panels
352         bpy.ops.object.editmode_toggle()
353         bpy.ops.object.editmode_toggle()
354         return {'FINISHED'}
355
356
357 class MESH_PT_ShowWeights(bpy.types.Panel):
358     bl_label = "Show Weights"
359     bl_space_type = "VIEW_3D"
360     bl_region_type = "UI"
361     bl_options = {'DEFAULT_CLOSED'}
362
363     @classmethod
364     def poll(cls, context):
365         return context.mode == 'EDIT_MESH'
366
367     def draw(self, context):
368         layout = self.layout
369         ob = context.active_object
370         me = ob.data
371         bm = bmesh.from_edit_mesh(me)
372         dvert_lay = bm.verts.layers.deform.active
373
374         row = layout.row(align=True)
375
376 #        text = "Show" if not context.scene.show_vgroups_weights else "Hide"
377 #        row.operator(ShowVGroupWeights.bl_idname, text=text)
378 #        row.prop(context.scene, "show_vgroups_weights_limit")
379
380         if len(ob.vertex_groups) > 0:
381             # Active vertex
382             active_vert = bm.select_history.active
383             sub = layout.box()
384             col = sub.column(align = True)
385             if bm.select_mode == {'VERT'} and active_vert is not None:
386                 col.label(text = "Active Vertex")
387                 row = col.row()
388                 row.label(text = "Vertex " + str(active_vert.index) + ":")
389                 row.operator_menu_enum("mesh.vertex_group_add", "available_vgroups", text = "Add Group", icon = 'GROUP_VERTEX')
390                 has_groups = False
391                 vgroup_weights = []
392
393                 for i in me.vertices:
394                     if i.index == active_vert.index:
395                         vgroup_weights_index = i.index
396                         for j in range(len(i.groups)):
397                             for k in ob.vertex_groups:
398                                 if k.index == i.groups[j].group:
399                                     has_groups = True
400                                     split = col.split(factor = 0.90, align = True)
401                                     vgroup_weights.append((k.index, i.groups[j].weight))
402                                     row = split.row(align = True)
403                                     row.prop(i.groups[j], "weight", text = k.name, slider = True, emboss = not k.lock_weight)
404                                     row = split.row(align = True)
405                                     row.operator("mesh.vertex_group_remove", text="", icon='X').vert_and_group = (i.index, k.index)
406
407                 if not has_groups:
408                     col.label(text = "    No Groups")
409                 else:
410                     props = col.operator("mesh.vertex_group_assign")
411                     props.index = vgroup_weights_index
412
413                     for vgroup, weight in vgroup_weights:
414                         item = props.vgroup_weights.add()
415                         item.vgroup = vgroup
416                         item.weight = weight
417
418                 layout.separator()
419             else:
420                 col.label(text = "No Active Vertex")
421             layout.prop(context.window_manager, "show_vgroups_show_all", toggle = True)
422             # All selected vertices (except for the active vertex)
423             if context.window_manager.show_vgroups_show_all:
424                 for v in bm.verts:
425                     if v.select:
426                         if active_vert is not None and v.index == active_vert.index:
427                             continue
428                         sub = layout.box()
429                         col = sub.column(align = True)
430                         col.label(text = "Vertex " + str(v.index) + ":")
431                         has_groups = False
432                         vgroup_weights = []
433                         for i in me.vertices:
434                             if i.index == v.index:
435                                 vgroup_weights_index = i.index
436                                 for j in range(len(i.groups)):
437                                     for k in ob.vertex_groups:
438                                         if k.index == i.groups[j].group:
439                                             has_groups = True
440                                             split = col.split(factor = 0.90, align = True)
441                                             vgroup_weights.append((k.index, i.groups[j].weight))
442                                             row = split.row(align = True)
443                                             row.prop(i.groups[j], "weight", text = k.name, slider = True, emboss = not k.lock_weight)
444                                             row = split.row(align = True)
445                                             row.operator("mesh.vertex_group_remove", text="", icon='X').vert_and_group = (i.index, k.index)
446                         if not has_groups:
447                             col.label(text = "    No Groups")
448                         else:
449                             props = col.operator("mesh.vertex_group_assign")
450                             props.index = vgroup_weights_index
451
452                             for vgroup, weight in vgroup_weights:
453                                 item = props.vgroup_weights.add()
454                                 item.vgroup = vgroup
455                                 item.weight = weight
456
457         else:
458             layout.label(text = "No Groups")
459
460
461 def create_properties():
462     bpy.types.WindowManager.show_vgroups_show_all = bpy.props.BoolProperty(
463         name = "Show All Selected Vertices",
464         description = "Show all vertices with vertex groups assigned to them",
465         default=False)
466
467     bpy.types.Mesh.assign_vgroup = bpy.props.StringProperty()
468
469     bpy.types.Scene.show_vgroups_weights = bpy.props.BoolProperty(
470         name="Show Vertex Groups/Weights",
471         default=False)
472
473     bpy.types.Scene.show_vgroups_weights_limit = bpy.props.IntProperty(
474         name="Limit",
475         description="Maximum number of weight overlays to draw",
476         default=20,
477         min=1,
478         max=1000,
479         soft_max=100)
480
481 # removal of ID-properties when script is disabled
482 def clear_properties(full=True):
483
484     if bpy.context.active_object is not None:
485         me = bpy.context.active_object.data
486
487         if hasattr(me, "show_vgroup_verts"):
488             del me["show_vgroup_verts"]
489         if hasattr(me, "show_vgroup_weights"):
490             del me["show_vgroup_weights"]
491
492     if full:
493         del bpy.types.WindowManager.show_vgroups_show_all
494         del bpy.types.Mesh.assign_vgroup
495         del bpy.types.Scene.show_vgroups_weights
496         del bpy.types.Scene.show_vgroups_weights_limit
497
498
499 classes = (
500     ShowVGroupWeights,
501     VGroupsWeights,
502     AssignVertexWeight,
503     RemoveFromVertexGroup,
504     AddToVertexGroup,
505     MESH_PT_ShowWeights
506 )
507
508 def register():
509     for cls in classes:
510         bpy.utils.register_class(cls)
511
512     create_properties()
513
514 def unregister():
515     ShowVGroupWeights.handle_remove()
516     clear_properties()
517
518     for cls in classes:
519         bpy.utils.unregister_class(cls)
520
521 if __name__ == "__main__":
522     register()