node_GreasePencil_In.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 #####


import bpy
from bpy.props import StringProperty, EnumProperty, BoolProperty
from bpy.props import IntProperty, FloatProperty
from node_s import *
from util import *
import ast
import os


FAIL_COLOR = (0.8, 0.1, 0.1)
READY_COLOR = (0, 0.8, 0.95)


def new_output_socket(node, name, stype):
    socket_type = {
        'v': 'VerticesSocket',
        's': 'StringsSocket',
        'm': 'MatrixSocket'
    }.get(stype, None)

    if socket_type:
        node.outputs.new(socket_type, name, name)


class SvGreaseNodeCallbackOp(bpy.types.Operator):

    bl_idname = "node.sverchok_greasenode_callback"
    bl_label = "Sverchok greasenode callback"
    bl_options = {'REGISTER', 'UNDO'}

    fn_name = StringProperty(default='')

    def execute(self, context):
        n = context.node
        f = getattr(n, self.fn_name, None)
        if f:
            f()
        return {'FINISHED'}


class SvGreasePencilNode(Node, SverchCustomTreeNode):

    ''' GreasePencil Input '''
    bl_idname = 'SvGreasePencilNode'
    bl_label = 'Grease Pencil Node'
    bl_icon = 'OUTLINER_OB_EMPTY'

    node_dict = {}  # this node's cache
    n_id = StringProperty(default='')

    frame_num = IntProperty(default=0, name='frame')
    frame_index = IntProperty(default=0, name='frame_idx')
    using_strokes = StringProperty(default="", update=updateNode)
    has_strokes = BoolProperty(default=False)

    close_distance = FloatProperty(default=1.0, update=updateNode)
    previous_distance = FloatProperty(default=1.0)  # manual update only

    def avail_layers(self, context):
        gp_data = bpy.data.grease_pencil
        items = []
        for gp in gp_data:
            for layer_name, layer_data in gp.layers.items():
                frames = layer_data.frames
                if not frames:
                    continue
                for frame_id, data in frames.items():

                    #if not (data.frame_number == self.frame_num):
                    #    continue

                    frame_num = 'frame: ' + str(data.frame_number)
                    ident = ' | '.join([gp.name, layer_name, frame_num, str(frame_id)])
                    # identifier, name, description
                    items.append((ident, ident, ''))
        return items

    layer_choices = EnumProperty(
        items=avail_layers,
        name="Layers in .blend",
        description="Choose layer to import and cache",
        update=updateNode)

    def init(self, context):
        self.n_id = node_id(self)
        self.node_dict[self.n_id] = {}  # or node_id(self) ?
        new_output_socket(self, 'Verts', 'v')
        new_output_socket(self, 'Edges', 's')

    def draw_buttons(self, context, layout):

        cb = 'node.sverchok_greasenode_callback'
        col = layout.column(align=True)
        if not self.has_strokes:
            #row = col.row(align=True)
            #row.prop(self, 'frame_num')
            row = col.row(align=True)
            row.label(text='load layer:')
            row = col.row(align=True)
            row.prop(self, 'layer_choices', '')
            row = col.row(align=True)
            row.operator(cb, text='Import Strokes').fn_name = 'load'
        else:
            row = col.row(align=True)
            row.prop(self, 'close_distance')
            row = col.row()
            label_info = 'loaded: ' + self.using_strokes
            col.label(icon='GREASEPENCIL', text=label_info)
            row = col.row(align=True)
            row.operator(cb, text='Load different').fn_name = 'load_new'

    def load(self):
        # print('self.layer_choices', self.layer_choices)
        self.node_dict[self.n_id] = {}

        self.using_strokes = self.layer_choices
        self.frame_index = int(self.layer_choices.split(' | ')[-1])
        self.has_strokes = True
        self.update()

    def load_new(self):
        self.using_strokes = ''
        self.has_strokes = False

        if 'current_cache' in self.node_dict[self.n_id]:
            del self.node_dict[self.n_id]['current_cache']

    def get_strokes(self):
        gp, layer_name, frame_n, frame_idx = self.using_strokes.split(' | ')
        gp = gp.strip()
        layer_name = layer_name.strip()

        # this is a little weird, frame_num and frame_index are not the same
        # frame idx is the entry in the list of frames it has nothing to do with time.
        # frame num is what frame this frame_idx is supposed to be on.

        # it's possible to arrive here when this data is invalidated.
        data = bpy.data.grease_pencil
        if not (gp in data and layer_name in data[gp].layers):
            return [], []

        layer = data[gp].layers[layer_name]
        strokes = layer.frames[self.frame_index].strokes

        all_verts = []
        all_edges = []
        for s in strokes:
            verts = []
            edges = []
            add_vert = verts.append
            for point in s.points:
                add_vert(point.co.to_tuple())

            all_verts.append(verts)

            # it's possible to have one point strokes.
            if len(verts) > 1:
                edges = [(n, n+1) for n in range(len(verts)-1)]

                # might want to make it cyclic (close the gap)
                if len(verts) >= 3:
                    if (s.points[0].co-s.points[-1].co).length <= self.close_distance:
                        edges.append((len(verts)-1, 0))

            else:
                # print('this stroke has 1 vertex')
                edges = [(0, 0)]
            all_edges.append(edges)

        return all_verts, all_edges

    def update(self):
        if not self.has_strokes and not self.using_strokes:
            return

        # use cache so long the close distance isn't changed
        unchanged = (self.previous_distance == self.close_distance)
        if ('current_cache' in self.node_dict[self.n_id]) and unchanged:
                [all_verts, all_edges] = self.node_dict[self.n_id]['current_cache']
        else:
            self.previous_distance = self.close_distance
            all_verts, all_edges = self.get_strokes()
            self.node_dict[self.n_id]['current_cache'] = [all_verts, all_edges]

        # possibly additional processing here ?

        SvSetSocketAnyType(self, 'Verts', all_verts)
        SvSetSocketAnyType(self, 'Edges', all_edges)


def register():
    bpy.utils.register_class(SvGreasePencilNode)
    bpy.utils.register_class(SvGreaseNodeCallbackOp)


def unregister():
    bpy.utils.unregister_class(SvGreasePencilNode)
    bpy.utils.unregister_class(SvGreaseNodeCallbackOp)