gp_interpolate/interpolate_strokes/operators.py

import bpy
import numpy as np
from time import perf_counter, time, sleep
from mathutils import Vector, Matrix

from ..utils import (matrix_transform,
                                  plane_on_bone,
                                  ray_cast_point,
                                  obj_ray_cast,
                                  intersect_with_tesselated_plane,
                                  triangle_normal,
                                  search_square,
                                  get_gp_draw_plane,
                                  create_plane,
                                  following_keys,
                                  index_list_from_bools,
                                  attr_set)

from mathutils.geometry import (barycentric_transform,
                                intersect_point_tri,
                                intersect_point_line,
                                intersect_line_plane,
                                tessellate_polygon)


class GP_OT_interpolate_stroke_base(bpy.types.Operator):
    bl_idname = "gp.interpolate_stroke_base"
    bl_label = "Interpolate Stroke"
    bl_description = 'Interpolate Stroke based on user bound triangle'
    bl_options = {'REGISTER', 'UNDO'} 

    next : bpy.props.BoolProperty(name='Next', default=True, options={'SKIP_SAVE'})

    @classmethod
    def poll(cls, context):
        if context.active_object and context.object.type == 'GPENCIL'\
            and context.mode in ('EDIT_GPENCIL', 'SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            return True
        cls.poll_message_set("Need a Grease pencil object in Edit or Sculpt mode")
        return False

    @classmethod
    def description(cls, context, properties):
        if properties.next:
            return f"Interpolate Stroke Forward"
        else:
            return f"Interpolate Stroke Backward"

    def apply_and_store(self):
        # self.store = []
        # item = (prop, attr, [new_val])
        for item in self.store_list:
            prop, attr = item[:2]
            self.store.append( (prop, attr, getattr(prop, attr)) )
            if len(item) >= 3:
                setattr(prop, attr, item[2])
    
    def restore(self):
        for prop, attr, old_val in self.store:
            setattr(prop, attr, old_val)

    def exit_modal(self, context, status='INFO', text=None, cancelled=False):
        context.area.header_text_set(None)
        wm = context.window_manager
        wm.progress_end() # Pgs
        wm.event_timer_remove(self.timer)
        self.restore()
        if self.debug:
            if self.scan_time is not None:
                print(f"Paste'n'place time {time()-self.start - self.scan_time}s")
        else:
            if self.settings.method == 'BONE':
                ## Remove Plane and it's collection after use
                if self.plane is not None:
                    bpy.data.objects.remove(self.plane)
                if self.toolcol is not None:
                    bpy.data.collections.remove(self.toolcol)

        cancel_state = '(Stopped!) ' if cancelled else ''
        mess = f'{cancel_state}{self.loop_count} interpolated frame(s) ({time()-self.start:.3f}s)'
        
        if text:
            print(mess)
            self.report({status}, text)
        else:
            self.report({'INFO'}, mess)

    def get_stroke_to_interpolate(self, context):
        if not self.gp.data.layers.active:
            self.exit_modal(context, status='ERROR', text='No active layer')
            return {'CANCELLED'}
        if not self.gp.data.layers.active.active_frame:
            self.exit_modal(context, status='ERROR', text='No active frame')
            return {'CANCELLED'}

        ## Get strokes to interpolate
        tgt_strokes = [s for s in self.gp.data.layers.active.active_frame.strokes if s.select]
        
        ## If nothing selected in sculpt/paint, Select all before triggering
        if not tgt_strokes and context.mode in ('SCULPT_GPENCIL', 'PAINT_GPENCIL'):
            for s in self.gp.data.layers.active.active_frame.strokes:
                s.select = True
            tgt_strokes = self.gp.data.layers.active.active_frame.strokes

        if not tgt_strokes:            
            self.exit_modal(context, status='ERROR', text='No stroke selected !')
            return {'CANCELLED'}
        return tgt_strokes

    ## For now, operators have their own invoke
    def invoke(self, context, event):
        self.debug = False
        self.store_list = []
        self.store = []
        self.loop_count = 0
        self.start = time()
        self.scan_time = None
        self.plane = None
        self.toolcol = None
        self.gp = context.object
        self.settings = context.scene.gp_interpo_settings
        self.frames_to_jump = None
        self.cancelled = False

        ## Remove interpolation_plane collection ! (unseen, but can be hit)
        if interp_plane := bpy.data.objects.get('interpolation_plane'):
            bpy.data.objects.remove(interp_plane)
        if interp_col := bpy.data.collections.get('interpolation_tool'):
            bpy.data.collections.remove(interp_col)

        ## Determine on what key/keys to jump
        self.frames_to_jump = following_keys(forward=self.next, animation=self.settings.use_animation)
        if not len(self.frames_to_jump):
            self.exit_modal(context, status='WARNING', text='No keyframe available in this direction')
            return {'CANCELLED'}
        
        self.timer = context.window_manager.event_timer_add(0.01, window=context.window)

## Converted to modal from "operator_single"

class GP_OT_interpolate_stroke(GP_OT_interpolate_stroke_base):
    bl_idname = "gp.interpolate_stroke"
    bl_label = "Interpolate Stroke"
    bl_description = 'Interpolate Stroke'
    bl_options = {'REGISTER', 'UNDO'} 

    def iterative_search(self, context, obj, coord, origin, dg):
        '''Search geometry for outside point (where raycast did not hit any geometry)
        
        return :
            object_hit, hit_location, tri, tri_indices.
            None if nothing found
        '''

        for iteration in range(1, 10):
            for square_co in search_square(coord, factor=self.settings.search_range * iteration):

                if obj:
                    object_hit, hit_location, tri, tri_indices = obj_ray_cast(obj, square_co, origin, dg)
                else:
                    object_hit, hit_location, tri, tri_indices = ray_cast_point(square_co, origin, dg)

                if object_hit:
                    ## On location coplanar with face triangle
                    # hit_location = intersect_line_plane(origin, coord, hit_location, triangle_normal(*tri))

                    ## On view plane 
                    view_vec = context.scene.camera.matrix_world.to_quaternion() @ Vector((0,0,1))
                    hit_location = intersect_line_plane(origin, coord, hit_location, view_vec)

                    ## An average of the two ?
                    # hit_location_1 = intersect_line_plane(origin, coord, hit_location, triangle_normal(*tri))
                    # hit_location_2 = intersect_line_plane(origin, coord, hit_location, view_vec)
                    # hit_location = (hit_location_1 + hit_location_2) / 2

                    return object_hit, hit_location, tri, tri_indices 

        return None, None, None, None

    def invoke(self, context, event):
        if state := super().invoke(context, event):
            return state

        scn = bpy.context.scene

        origin = scn.camera.matrix_world.to_translation()
        
        tgt_strokes = self.get_stroke_to_interpolate(context)
        if isinstance(tgt_strokes, set):
            return tgt_strokes

        col = self.settings.target_collection
        if not col:
            col = scn.collection

        included_cols = [c.name for c in self.gp.users_collection]
        target_obj = None
        
        ## Setup depending on method
        if self.settings.method == 'BONE':
            if not self.settings.target_rig or not self.settings.target_bone:
                self.exit_modal(context, status='ERROR', text='No Bone selected')
                return {'CANCELLED'}

            included_cols.append('interpolation_tool')

            ## Ensure collection and plane exists
            # get/create collection
            self.toolcol = bpy.data.collections.get('interpolation_tool')
            if not self.toolcol:
                self.toolcol = bpy.data.collections.new('interpolation_tool')

            if self.toolcol.name not in bpy.context.scene.collection.children:
                bpy.context.scene.collection.children.link(self.toolcol)
                self.toolcol.hide_viewport = True # needed ?
            
            # get/create meshplane
            self.plane = bpy.data.objects.get('interpolation_plane')
            if not self.plane:
                self.plane = create_plane(name='interpolation_plane')
                self.plane.select_set(False)

            if self.plane.name not in self.toolcol.objects:
                self.toolcol.objects.link(self.plane)
            target_obj = self.plane

        elif self.settings.method == 'GEOMETRY':
            if col != context.scene.collection:
                included_cols.append(col.name)
        
        elif self.settings.method == 'OBJECT':
            if not self.settings.target_object:
                self.exit_modal(context, status='ERROR', text='No Object selected')
                return {'CANCELLED'}
            col = scn.collection # Reset collection filter
            target_obj = self.settings.target_object
            if target_obj.library:
                ## Look if an override exists in scene to use instead of default object
                if (override := next((o for o in scn.objects if o.name == target_obj.name and o.override_library), None)):
                    target_obj = override

        ## Prepare context manager
        self.store_list = [
            # (context.view_layer.objects, 'active', self.gp),
            (context.tool_settings, 'use_keyframe_insert_auto', True),
            # (bpy.context.scene.render, 'simplify_subdivision', 0),
            ]

        ## Set everything in SETUP list
        self.apply_and_store()

        if self.settings.method == 'BONE':
            ## replace plane
            _bone_plane = plane_on_bone(self.settings.target_rig.pose.bones.get(self.settings.target_bone),
                                        arm=self.settings.target_rig,
                                        set_rotation=self.settings.use_bone_rotation,
                                        mesh=True)

            ## Set collection visibility
            intercol = bpy.data.collections.get('interpolation_tool')
            vl_col = bpy.context.view_layer.layer_collection.children.get(intercol.name)
            intercol.hide_viewport = vl_col.exclude = vl_col.hide_viewport = False

            ## Override collection
            col = intercol
        
        dg = bpy.context.evaluated_depsgraph_get()
        self.strokes_data = []

        for stroke_index, stroke in enumerate(tgt_strokes):
            stroke_data = []
            for point_index, point in enumerate(stroke.points):
                point_co_world = self.gp.matrix_world @ point.co

                if target_obj:
                    ## Object raycast
                    object_hit, hit_location, tri, tri_indices = obj_ray_cast(target_obj, point_co_world, origin, dg)
                else:
                    ## Scene raycast
                    object_hit, hit_location, tri, tri_indices = ray_cast_point(point_co_world, origin, dg)

                ## Iterative increasing search range when no surface hit
                if not object_hit:
                    object_hit, hit_location, tri, tri_indices = self.iterative_search(context, target_obj, point_co_world, origin, dg)

                    if not object_hit:
                        ## /!\ ERROR ! No surface found!
                        # For debugging, select only point.
                        bpy.ops.gpencil.select_all(action='DESELECT')
                        point.select = True
                        self.exit_modal(context, status='ERROR', text=f'Stroke {stroke_index} point {point_index} could not find underlying geometry')
                        return {'CANCELLED'}
                
                stroke_data.append((stroke, point_co_world, object_hit, hit_location, tri, tri_indices))

            self.strokes_data.append(stroke_data)

        if self.debug:
            self.scan_time = time()-self.start
            print(f'Scan time {self.scan_time:.4f}s')
        
        # Copy stroke selection
        bpy.ops.gpencil.select_linked() # Ensure whole stroke are selected before copy
        bpy.ops.gpencil.copy()

        # Jump frame and paste

        context.window_manager.progress_begin(self.frames_to_jump[0], self.frames_to_jump[-1]) # Pgs
        context.window_manager.modal_handler_add(self)
        context.area.header_text_set('Starting interpolation | Esc: Cancel')

        return {'RUNNING_MODAL'}

    def modal(self, context, event):
        frame_num = len(self.frames_to_jump)
        percentage = (self.loop_count) / (frame_num) * 100
        context.area.header_text_set(f'Interpolation {percentage:.0f}% {self.loop_count + 1}/{frame_num} | Esc: Cancel')
        # (frame: {self.frames_to_jump[self.loop_count]})

        if event.type in {'RIGHTMOUSE', 'ESC'}:
            context.area.header_text_set(f'Cancelling')
            self.exit_modal(context, text='Cancelling', cancelled=True)
            return {'CANCELLED'}

        ## -- LOOPTIMER 
        if event.type == 'TIMER':
            f = self.frames_to_jump[self.loop_count]
            bpy.context.window_manager.progress_update(f) # Pgs
            scn = bpy.context.scene
            scn.frame_set(f)
            origin = scn.camera.matrix_world.to_translation()
            plane_co, plane_no = get_gp_draw_plane(self.gp)
            bpy.ops.gpencil.select_all(action='DESELECT')
            bpy.ops.gpencil.paste()

            if self.settings.method == 'BONE':
                ## Set plane on the bone
                plane_on_bone(self.settings.target_rig.pose.bones.get(self.settings.target_bone),
                                            arm=self.settings.target_rig,
                                            set_rotation=self.settings.use_bone_rotation,
                                            mesh=True)

            dg = bpy.context.evaluated_depsgraph_get()
            
            ## Get pasted stroke
            new_strokes = [s for s in self.gp.data.layers.active.active_frame.strokes if s.select]
            ## Keep reference to all accessible other strokes (in all accessible layer)
            other_strokes = [s for l in self.gp.data.layers if l.active_frame and not l.lock for s in l.active_frame.strokes if not s.select]
            
            occluded_points = []
            # For new_stroke, stroke_data in zip(new_strokes, self.strokes_data):
            for new_stroke, stroke_data in zip(list(new_strokes), list(self.strokes_data)):
                world_co_3d = []
                for stroke, point_co, object_hit, hit_location, tri_a, tri_indices in stroke_data:
                    eval_ob = object_hit.evaluated_get(dg)
                    tri_b = [eval_ob.matrix_world @ eval_ob.data.vertices[i].co for i in tri_indices]

                    new_loc = barycentric_transform(hit_location, *tri_a, *tri_b)                        
                    world_co_3d.append(new_loc)

                ## Reproject on plane
                new_world_co_3d = [intersect_line_plane(origin, p, plane_co, plane_no) for p in world_co_3d]        
                new_local_co_3d = [co for coord in new_world_co_3d for co in self.gp.matrix_world.inverted() @ coord]
                new_stroke.points.foreach_set('co', new_local_co_3d)
                new_stroke.points.update()

                ## Occlusion management
                if self.settings.method == 'GEOMETRY' and self.settings.remove_occluded:
                    for i, point in enumerate(new_stroke.points):
                        point_co = world_co_3d[i]
                        vec_direction = point_co - origin
                        ## Raycast with slightly reduced distance (avoid occlusion on same source)
                        n_hit, _, _, _, _, _ = scn.ray_cast(dg, origin, vec_direction, distance=vec_direction.length - 0.001)
                        if n_hit:
                            occluded_points.append(point)

            if occluded_points:
                ## Select only occluded point
                bpy.ops.gpencil.select_all(action='DESELECT')
                for point in occluded_points:
                    point.select = True
                ## remove points
                bpy.ops.gpencil.delete(type='POINTS')

                ## restore selection (keep new strokes selected)
                bpy.ops.gpencil.select_all(action='SELECT')
                for stroke in other_strokes:
                    stroke.select = False
                    
            self.loop_count += 1
            if self.loop_count >= len(self.frames_to_jump):
                self.exit_modal(context)
                return {'FINISHED'}
            
            bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
            # context.area.tag_redraw()

        return {'RUNNING_MODAL'}


classes = (
    GP_OT_interpolate_stroke_base,
    GP_OT_interpolate_stroke,
)

def register():
    for c in classes:
        bpy.utils.register_class(c)
    
        
def unregister():
    for c in reversed(classes):
        bpy.utils.unregister_class(c)