Vertex Color Helper (RAWR!)

view3d_vertex_color_view.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  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, write to the Free Software Foundation,
#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####


bl_info = {
    'name': 'Vertex Color Visualiser',
    'author': 'zeffii',
    'version': (0, 0, 1),
    'blender': (2, 7, 4),
    'location': 'View3D > Properties panel > Mesh Display tab (edit-mode)',
    'warning': '',
    'description': 'Display indices of verts, edges and faces in the 3d-view',
    'wiki_url': '',
    'tracker_url': '',
    'category': '3D View'}


"""
Some of this code borrows heavily from IndexVisualizer.

"""

import bpy
import bgl
import blf
import mathutils
import bmesh
import math
from bpy.props import IntProperty, BoolProperty, FloatVectorProperty
from bpy_extras.view3d_utils import location_3d_to_region_2d as loc3d2d

Scene = bpy.types.Scene

point_dict = {}


def adjust_list(in_list, x, y):
    return [[old_x + x, old_y + y] for (old_x, old_y) in in_list]


def generate_points(width, height):
    amp = 5  # radius fillet

    width += 2
    height += 4
    width = ((width / 2) - amp) + 2
    height -= (2 * amp)

    pos_list, final_list = [], []

    n_points = 12
    seg_angle = 2 * math.pi / n_points
    for i in range(n_points + 1):
        angle = i * seg_angle
        x = math.cos(angle) * amp
        y = math.sin(angle) * amp
        pos_list.append([x, -y])

    w_list, h_list = [1, -1, -1, 1], [-1, -1, 1, 1]
    slice_list = [[i, i + 4] for i in range(0, n_points, 3)]

    for idx, (start, end) in enumerate(slice_list):
        point_array = pos_list[start:end]
        w = width * w_list[idx]
        h = height * h_list[idx]
        final_list += adjust_list(point_array, w, h)

    return final_list


def get_points(index):
    '''
    index:   string representation of the index number
    returns: rounded rect point_list used for background.

    the neat thing about this is if a width has been calculated once, it
    is stored in a dict and used if another polygon is saught with that width.
    '''
    width, height = blf.dimensions(0, index)
    if not (width in point_dict):
        point_dict[width] = generate_points(width, height)

    return point_dict[width]


def vcol_info(ob, map_name, poly_idx):
    color_map_collection = ob.data.vertex_colors
    color_map = color_map_collection[map_name]

    i = 0
    mesh = ob.data
    for poly in mesh.polygons:
        for idx in poly.loop_indices:
            loop = mesh.loops[idx]
            vertex_index = loop.vertex_index
            color_map.data[i].color = val
            i += 1


def draw_callback_px(self, context):
    if context.mode != "PAINT_VERTEX":
        return

    # get screen information
    region = context.region
    rv3d = context.space_data.region_3d
    this_object = context.active_object
    matrix_world = this_object.matrix_world

    text_height = 13
    blf.size(0, text_height, 72)

    def draw_index(rgb, index, coord):

        vector3d = matrix_world * coord
        x, y = loc3d2d(region, rv3d, vector3d)

        index = str(index)
        polyline = get_points(index)

        ''' draw polygon '''
        bgl.glColor4f(0.103, 0.2, 0.2, 0.2)
        bgl.glBegin(bgl.GL_POLYGON)
        for pointx, pointy in polyline:
            bgl.glVertex2f(pointx + x, pointy + y)
        bgl.glEnd()

        ''' draw text '''
        txt_width, txt_height = blf.dimensions(0, index)
        bgl.glColor3f(*rgb)
        blf.position(0, x - (txt_width / 2), y - (txt_height / 2), 0)
        blf.draw(0, index)

    bgl.glColor3f(0.8, 0.2, 0.9)
    blf.position(0, 30, 30, 0)  # fontid, x, y, z
    blf.draw(0, "monkies!")

    vert_idx_color = (1.0, 1.0, 1.0)
    face_idx_color = (1.0, 0.8, 0.8)

    scene = context.scene
    me = context.active_object.data

    # Get a BMesh representation
    bm = bmesh.new()   # create an empty BMesh
    bm.from_mesh(me)   # fill it in from a Mesh

    if scene.live_mode:
        me.update()

    if scene.display_vert_index:
        for v in bm.verts:
            if not v.hide and (v.select or not scene.display_sel_only):
                draw_index(vert_idx_color, v.index, v.co.to_4d())

    if scene.display_face_index:
        for f in bm.faces:
            if not f.hide and (f.select or not scene.display_sel_only):
                draw_index(
                    face_idx_color, f.index, f.calc_center_median().to_4d())


# operator
class VCOLHelper(bpy.types.Operator):
    bl_idname = "view3d.vcol_helper"
    bl_label = "Vertex Color Helper"
    bl_description = "Show and set the Vertex Colors per face"

    _handle = None

    @classmethod
    def poll(cls, context):
        return context.mode == "PAINT_VERTEX"

    def modal(self, context, event):
        if context.area:
            context.area.tag_redraw()

        # removal of callbacks when operator is called again
        if context.scene.display_indices == -1:
            bpy.types.SpaceView3D.draw_handler_remove(self._handle, 'WINDOW')
            context.scene.display_indices = 0
            return {"CANCELLED"}

        return {"PASS_THROUGH"}

    def invoke(self, context, event):
        if context.area.type == "VIEW_3D":
            if context.scene.display_indices < 1:
                # operator is called for the first time, start everything
                context.scene.display_indices = 1
                self._handle = bpy.types.SpaceView3D.draw_handler_add(
                    draw_callback_px, (self, context), 'WINDOW', 'POST_PIXEL')
                context.window_manager.modal_handler_add(self)
                return {"RUNNING_MODAL"}
            else:
                # operator is called again, stop displaying
                context.scene.display_indices = -1
                return {'RUNNING_MODAL'}
        else:
            self.report({"WARNING"}, "View3D not found, can't run operator")
            return {"CANCELLED"}


# defining the panel
def menu_func(self, context):
    self.layout.separator()
    scn = context.scene

    col = self.layout.column(align=True)
    col.operator(VCOLHelper.bl_idname, text="Get / Set VertexColors")
    row = col.row(align=True)
    row.active = (context.mode == "PAINT_VERTEX" and scn.display_indices == 1)

    row.prop(scn, "display_vert_index", toggle=True)
    row.prop(scn, "display_face_index", toggle=True)
    row = col.row(align=True)
    row.active = context.mode == "PAINT_VERTEX" and scn.display_indices == 1
    row.prop(scn, "display_sel_only")

    col = self.layout.column()
    col.prop(scn, "vcol_helper_color_1")
    col.prop(scn, "vcol_helper_color_2")
    col.prop(scn, "vcol_helper_color_3")
    col.prop(scn, "vcol_helper_color_4")


def properties():
    def get_color_prop():
        WHITE = (1, 1, 1)
        return FloatVectorProperty(default=WHITE, subtype='COLOR', min=0.0, max=1.0)

    props = dict(
        display_indices=IntProperty(
            name="Display indices", default=0),
        display_sel_only=BoolProperty(
            name="Selected only",
            description="Only display indices of selected vertices/edges/faces",
            default=True),
        display_vert_index=BoolProperty(
            name="Vertices",
            description="Display vertex indices",
            default=True),
        display_face_index=BoolProperty(
            name="Faces",
            description="Display face indices"),
        live_mode=BoolProperty(
            name="Live",
            description="Toggle live update of the selection, can be slow",
            default=False),
        vcol_helper_color_1=get_color_prop(),
        vcol_helper_color_2=get_color_prop(),
        vcol_helper_color_3=get_color_prop(),
        vcol_helper_color_4=get_color_prop()
    )
    return props


def register_properties():
    for prop_name, prop_config in properties().items():
        setattr(Scene, prop_name, prop_config)


def unregister_properties():
    for prop_name in properties().keys():
        exec('del Scene.{0}'.format(prop_name))


def register():
    register_properties()
    bpy.utils.register_class(VCOLHelper)
    bpy.types.VIEW3D_PT_tools_brush.append(menu_func)


def unregister():
    bpy.utils.unregister_class(VCOLHelper)
    bpy.types.VIEW3D_PT_tools_brush.remove(menu_func)
    unregister_properties()


if __name__ == "__main__":
    register()