gp_interpolate/interpolate_strokes/operators_triangle.py

import bpy
from time import time

from mathutils.geometry import (barycentric_transform,
                                intersect_line_plane)

from ..utils import (triangle_normal,
                    get_gp_draw_plane)

from .operators import GP_OT_interpolate_stroke_base


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

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

        if not context.window_manager.get(f'tri_{self.gp.name}'):
            return self.exit(context, status='ERROR', text='Need to bind coordinate first. Use "Bind Tri Point" button')

        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

        ## Prepare context manager
        attrs = [
            # (context.view_layer.objects, 'active', self.gp),
            (context.tool_settings, 'use_keyframe_insert_auto', True),
            # (bpy.context.scene.render, 'simplify_subdivision', 0),
            ]
        self.apply_and_store(attrs)
        
        point_dict = context.window_manager.get(f'tri_{self.gp.name}')
        ## point_dict -> {'0': {'object': object_name_as_str, 'index': 450}, ...}
        ## Get triangle dumped in context.window_manager
        self.source_object_list = [bpy.context.scene.objects.get(point_dict[str(i)]['object']) for i in range(3)]
        self.source_tri_indices = [point_dict[str(i)]['index'] for i in range(3)] # List of vertices index corresponding to tri coordinates
        
        dg = bpy.context.evaluated_depsgraph_get()
        
        ## Get tri at source frame
        tri = []
        for source_obj, idx in zip(self.source_object_list, self.source_tri_indices):
            ob_eval = source_obj.evaluated_get(dg)
            tri.append(ob_eval.matrix_world @ ob_eval.data.vertices[idx].co)

        self.strokes_data = []

        for stroke in tgt_strokes:
            stroke_data = []
            for point in stroke.points:
                point_co_world = self.gp.matrix_world @ point.co

                ## Set hit location at same coordinate as point
                # hit_location = point_co_world
                
                ## Set hit location on tri plane
                hit_location = intersect_line_plane(origin, point_co_world, tri[0], triangle_normal(*tri))

                stroke_data.append((hit_location, tri))

            self.strokes_data.append(stroke_data)

        if self.debug:
            self.scan_time = time()-self.start
            print(f'Scan time {self.scan_time:.4f}s')

        # Ensure whole stroke are selected before copy
        bpy.ops.gpencil.select_linked()
        # Copy stroke selection
        bpy.ops.gpencil.copy()

        # Jump frame and paste
        # if self.report_progress:
        #     context.window_manager.progress_begin(self.frames_to_jump[0], self.frames_to_jump[-1]) # Pgs
        # context.area.header_text_set('Starting interpolation | Esc: Cancel')

        context.window_manager.modal_handler_add(self)
        return {'RUNNING_MODAL'}

    def interpolate_frame(self, context):
        scn = context.scene
        origin = scn.camera.matrix_world.to_translation()
        plane_co, plane_no = get_gp_draw_plane(self.gp)
        bpy.ops.gpencil.paste()

        dg = bpy.context.evaluated_depsgraph_get()
        
        ## List of newly pasted strokes (using range)
        new_strokes = self.gp.data.layers.active.active_frame.strokes[-len(self.strokes_data):]

        ## Get user triangle position at current frame
        tri_b = []
        for source_obj, idx in zip(self.source_object_list, self.source_tri_indices):
            ob_eval = source_obj.evaluated_get(dg)
            tri_b.append(ob_eval.matrix_world @ ob_eval.data.vertices[idx].co)

        for new_stroke, stroke_data in zip(list(new_strokes), list(self.strokes_data)):
            world_co_3d = []
            for hit_location, tri_a in stroke_data:
                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()

classes = (
    GP_OT_interpolate_stroke_tri,
)

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