space_view_3d_mesh_vertex_marker.py

#
# Author: Andrei Izrantcev
#
# With the great help of https://github.com/zeffii/ repository
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

bl_info = {
    'name'        : "Mesh Vertex Marker",
    'author'      : "Andrei Izrantcev <izrantsev@gmail.com>",
    'version'     : (0, 0, 1),
    'blender'     : (2, 5, 8),
    'location'    : "3d view > Tools > Misc > Mesh Vertex Marker",
    'description' : "Marks vertices with 3 or 5 edges",
    'wiki_url'    : "",
    'tracker_url' : "",
    'category'    : "Mesh"
}


import bpy
import bgl
import bmesh
import mathutils
import bpy_extras

from bpy_extras.view3d_utils import location_3d_to_region_2d


def mesh_vertex_marker_draw_handler(self, context):
    ob = context.object
    if not ob:
        return
    if ob.type not in {'MESH'}:
        return
    if ob.mode not in {'EDIT'}:
        return

    pg = context.scene.mesh_vertex_marker

    MarkColors = {
        3 : pg.color_3,
        # 4 : (0.0, 1.0, 0.0, 1.0),
        5 : pg.color_5,
    }

    bm = bmesh.from_edit_mesh(ob.data)
    if hasattr(bm.verts, 'ensure_lookup_table'):
        bm.verts.ensure_lookup_table()
        bm.edges.ensure_lookup_table()

    vertex_usage = [0] * len(bm.verts)
    for e in bm.edges:
        if e.hide:
            continue
        if pg.ignore_borders and len(e.link_faces) == 1:
            continue
        for v in e.verts:
            vertex_usage[v.index] += 1

    bgl.glEnable(bgl.GL_POINT_SMOOTH)
    bgl.glPointSize(pg.point_size)
    bgl.glBlendFunc(bgl.GL_SRC_ALPHA, bgl.GL_ONE_MINUS_SRC_ALPHA)

    bgl.glBegin(bgl.GL_POINTS)
    for vi,vu in enumerate(vertex_usage):
        if vu not in MarkColors:
            continue

        gl_col = MarkColors[vu]
        coord  = bm.verts[vi]

        vector3d = ob.matrix_world * mathutils.Vector((coord.co[0], coord.co[1], coord.co[2]))

        vector2d = location_3d_to_region_2d(context.region, context.space_data.region_3d, vector3d)

        bgl.glColor4f(*gl_col)
        bgl.glVertex2f(*vector2d)
    bgl.glEnd()

    bgl.glDisable(bgl.GL_POINTS)
    bgl.glDisable(bgl.GL_POINT_SMOOTH)

    bgl.glLineWidth(1)
    bgl.glDisable(bgl.GL_BLEND)
    bgl.glColor4f(0.0, 0.0, 0.0, 1.0)


class MeshVertexMarkerUI(bpy.types.Panel):
    bl_label       = "Mesh Vertex Marker"
    bl_space_type  = 'VIEW_3D'
    bl_region_type = 'TOOLS'
    bl_category    = 'Misc'

    @classmethod
    def poll(self, context):
        ob = context.object
        if ob and ob.mode == 'EDIT':
            return True

    def draw(self, context):
        pg = context.scene.mesh_vertex_marker

        box = self.layout.box()
        split = box.split(percentage=0.2)
        row = split.column()
        row.label("3:")
        row = split.column()
        row.prop(pg, 'color_3', text="")

        split = box.split(percentage=0.2)
        row = split.column()
        row.label("5:")
        row = split.column()
        row.prop(pg, 'color_5', text="")

        box.prop(pg, 'point_size')
        box.prop(pg, 'ignore_borders')

        op_icon = 'OUTLINER_OB_MESH' if MeshVertexMarkerOp._handle else 'MESH_DATA'

        self.layout.operator("mesh.vertex_maker", icon=op_icon)


class MeshVertexMarkerOp(bpy.types.Operator):
    bl_idname      = "mesh.vertex_maker"
    bl_label       = "Mesh Vertex Marker"
    bl_description = "Marks 3-5-edges vertices"

    _handle = None

    @classmethod
    def handle_add(cls, slf, ctx):
        if MeshVertexMarkerOp._handle:
            MeshVertexMarkerOp.handle_remove()

        MeshVertexMarkerOp._handle = bpy.types.SpaceView3D.draw_handler_add(
            mesh_vertex_marker_draw_handler,
            (slf, ctx),
            'WINDOW',
            'POST_PIXEL'
        )

    @classmethod
    def handle_remove(cls):
        if MeshVertexMarkerOp._handle:
            bpy.types.SpaceView3D.draw_handler_remove(
                MeshVertexMarkerOp._handle,
                'WINDOW'
            )
        MeshVertexMarkerOp._handle = None

    def execute(self, context):
        if MeshVertexMarkerOp._handle:
            MeshVertexMarkerOp.handle_remove()
        else:
            MeshVertexMarkerOp.handle_add(self, context)
        if context.area:
            context.area.tag_redraw()
        return {'FINISHED'}


class MeshVertexMarkerPg(bpy.types.PropertyGroup):
    color_3 = bpy.props.FloatVectorProperty(
        name        = "Color 3",
        description = "Color for 3 edges vertex",
        subtype     = 'COLOR',
        min         = 0.0,
        max         = 1.0,
        size        = 4,
        default     = (1.0,0.0,0.0,1.0)
    )

    color_5 = bpy.props.FloatVectorProperty(
        name        = "Color 5",
        description = "Color for 5 edges vertex",
        subtype     = 'COLOR',
        min         = 0.0,
        max         = 1.0,
        size        = 4,
        default     = (0.0,0.0,1.0,1.0)
    )

    point_size = bpy.props.FloatProperty(
        name        = "Size",
        description = "Draw point size",
        min         = 1.0,
        default     = 7.0
    )

    ignore_borders = bpy.props.BoolProperty(
        name        = "Ignore Borders",
        description = "Ignore border edges",
        default     = False
    )


RegClasses = (
    MeshVertexMarkerOp,
    MeshVertexMarkerUI,
    MeshVertexMarkerPg,
)


def register():
    for classname in RegClasses:
        bpy.utils.register_class(classname)

    bpy.types.Scene.mesh_vertex_marker = bpy.props.PointerProperty(
        name = "Mesh Vertex Marker",
        type =  bpy.types.MeshVertexMarkerPg,
        description = "Mesh vertex marker settings"
    )


def unregister():
    for classname in RegClasses:
        bpy.utils.unregister_class(classname)
    del bpy.types.Scene.mesh_vertex_marker


if __name__ == "__main__":
    register()