gp_interpolate/utils.py

import bpy
import math

import numpy as np

from math import tan
from mathutils import Vector, Matrix
from bpy_extras.object_utils import world_to_camera_view
from mathutils.geometry import (barycentric_transform, intersect_point_tri, 
    intersect_point_line, intersect_line_plane, tessellate_polygon)

## context manager to store restore

class attr_set():
    '''Receive a list of tuple [(data_path, "attribute" [, wanted value)] ]
    entering with-statement : Store existing values, assign wanted value (if any)
    exiting with-statement: Restore values to their old values
    '''

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

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        for prop, attr, old_val in self.store:
            setattr(prop, attr, old_val)

# --- Vector

def triangle_normal(a, b, c):
    x = a[1] * b[2] - a[2] * b[1]
    y = a[2] * b[0] - a[0] * b[2]
    z = a[0] * b[1] - a[1] * b[0]
    
    return np.array([x, y, z], dtype='float64')

def plane_coords(size=1):
    v = size * 0.5
    return np.array([(-v, v, 0), (v, v, 0), (v, -v, 0), (-v, -v, 0)], dtype='float64')

def matrix_transform(coords, matrix):
    coords_4d = np.column_stack((coords, np.ones(len(coords), dtype='float64')))
    return np.einsum('ij,aj->ai', matrix, coords_4d)[:, :-1]

def vector_normalized(vec):
    return vec / np.sqrt(np.sum(vec**2))

def vector_magnitude(vec):
    return np.sqrt(vec.dot(vec))

def search_square(point, factor=0.05, cam=None):
    if cam is None:
        cam = bpy.context.scene.camera
        
    plane = plane_coords()
    mat = cam.matrix_world.copy()
    mat.translation = point
    depth = vector_magnitude(point - cam.matrix_world.to_translation())
    mat_scale = Matrix.Scale(tan(cam.data.angle * 0.5) * depth * factor, 4)
    
    return matrix_transform(plane, mat @ mat_scale)

def get_tri_from_face(hit_location, face_index, object_hit, depsgraph):
    eval_ob = object_hit.evaluated_get(depsgraph)
    face = eval_ob.data.polygons[face_index]
    vertices = [eval_ob.data.vertices[i] for i in face.vertices]
    face_co = matrix_transform([v.co for v in vertices], eval_ob.matrix_world)
    
    tri = None
    for tri_idx in tessellate_polygon([face_co]):
        tri = [face_co[i] for i in tri_idx]
        tri_indices = [vertices[i].index for i in tri_idx]
        if intersect_point_tri(hit_location, *tri):
            break
    
    return tri, tri_indices

def ray_cast_point(point, origin, depsgraph):
    ray = (point - origin)
    hit, hit_location, normal, face_index, object_hit, matrix = bpy.context.scene.ray_cast(depsgraph, origin, ray)
    
    if not hit:
        return None, None, None, None
    
    tri, tri_indices = get_tri_from_face(hit_location, face_index, object_hit, depsgraph)

    return object_hit, np.array(hit_location), tri, tri_indices

def obj_ray_cast(obj, point, origin, depsgraph):
    """Wrapper for ray casting that moves the ray into object space"""

    # get the ray relative to the object
    matrix_inv = obj.matrix_world.inverted()
    ray_origin_obj = matrix_inv @ origin # matrix_transform(origin, matrix_inv) 
    ray_target_obj = matrix_inv @ point # matrix_transform(point, matrix_inv) 
    ray_direction_obj = ray_target_obj - ray_origin_obj
    # cast the ray
    success, location, normal, face_index = obj.ray_cast(ray_origin_obj, ray_direction_obj, depsgraph=depsgraph)
    if not success:
        return None, None, None, None

    # Get hit location world_space
    hit_location = obj.matrix_world @ location
    tri, tri_indices = get_tri_from_face(hit_location, face_index, obj, depsgraph)
    return obj, np.array(hit_location), tri, tri_indices


def empty_at(name='Empty', pos=(0,0,0), collection=None, type='PLAIN_AXES', size=1, show_name=False):
    '''
    Create an empty at given Vector3 position.
    Optional type (default 'PLAIN_AXES') in ,'ARROWS','SINGLE_ARROW','CIRCLE','CUBE','SPHERE','CONE','IMAGE'
    default size is 1.0
    '''
    
    mire = bpy.data.objects.get(name)
    if not mire:
        mire = bpy.data.objects.new(name, None)
    if collection is None:
        collection = bpy.context.collection
    if mire.name not in collection.all_objects:
        collection.objects.link(mire)

    mire.location = pos
    mire.empty_display_type = type
    mire.empty_display_size = size
    mire.show_name = show_name
    return mire

def plane_on_bone(bone, arm=None, cam=None, set_rotation=True, mesh=True):
    '''
    bone (posebone): reference pose bone
    arm (optional: Armature): Armature of the pose bone (if not passed found using bone.id_data)
    cam (optional: Camera) : Camera to align plane to (if not passed use scene camera)
    set_rotation (bool): rotate the plane on cam view axis according to bone direction in 2d cam space
    mesh (bool): create a real mesh ans return it, else return list of plane coordinate
    '''

    if cam is None:
        cam = bpy.context.scene.camera
    
    if arm is None:
        arm = bone.id_data

    mat = cam.matrix_world.copy()

    if set_rotation:
        head_world_coord = arm.matrix_world @ bone.head
        mat.translation = head_world_coord

        ## Apply 2d bone rotation facing camera

        # Get 2d camera space coords (NDC: normalized device coordinate, 0,0 is bottom-left)
        head_2d, tail_2d = get_bone_head_tail_2d(bone, cam=cam)
        vec_from_corner_2d = (tail_2d - head_2d).normalized()
        up_vec_2d = Vector((0,1))
        # angle = acos(up_vec_2d.dot(vec_from_corner_2d)) ## equivalent but not signed!
        angle = up_vec_2d.angle_signed(vec_from_corner_2d)

        ## Axis camera aim (seem slightly off)
        # rot_axis = Vector((0, 0, -1))
        # rot_axis.rotate(cam.matrix_world)

        ## Axis camera origin -> pivot
        rot_axis = head_world_coord - cam.matrix_world.translation
        
        mat = rotate_matrix_around_pivot(mat, angle, head_world_coord, rot_axis)

    else:
        ## Use mid bone to better follow movement
        mat.translation = arm.matrix_world @ ((bone.tail + bone.head) / 2) # Mid bone
    

    if mesh:
        # get/create collection
        col = bpy.data.collections.get('interpolation_tool')
        if not col:
            col = bpy.data.collections.new('interpolation_tool')

        if col.name not in bpy.context.scene.collection.children:
            bpy.context.scene.collection.children.link(col)
        
        # get/create meshplane
        plane = bpy.data.objects.get('interpolation_plane')
        if not plane:
            plane = create_plane(name='interpolation_plane')
            # Display type as Wire for a discrete XP
        plane.display_type = 'WIRE'

        if plane.name not in col.objects:
            col.objects.link(plane)

        plane.matrix_world = mat
        return plane

    mat_scale = Matrix.Scale(10, 4)
    plane = plane_coords()
    return matrix_transform(plane, mat @ mat_scale)

def place_object_to_ref_facing_cam(obj, ref_ob, bone=None, cam=None, set_rotation=True):
    '''
    obj (Object): the object to place
    ref_ob (Object): the reference object or armature
    bone (posebone): reference pose bone
    arm (optional: Armature): Armature of the pose bone (if not passed found using bone.id_data)
    cam (optional: Camera) : Camera to align plane to (if not passed use scene camera)
    set_rotation (bool): rotate the plane on cam view axis according to bone direction in 2d cam space
    '''

    if cam is None:
        cam = bpy.context.scene.camera
    
    # if ref_ob is None:
    #     ref_ob = bone.id_data

    mat = cam.matrix_world.copy()

    if set_rotation:
        if bone:
            head_world_coord = ref_ob.matrix_world @ bone.head
            mat.translation = head_world_coord

            ## Apply 2d bone rotation facing camera
            # Get 2d camera space coords (NDC: normalized device coordinate, 0,0 is bottom-left)
            head_2d, tail_2d = get_bone_head_tail_2d(bone, cam=cam)
        else:
            mat.translation = ref_ob.matrix_world
            # Get 2d camera space coords (NDC: normalized device coordinate, 0,0 is bottom-left)
            scene = bpy.context.scene
            up_vec = Vector((0,0,1))
            up_vec.rotate(ref_ob.matrix_world)
            tail_3d = ref_ob.matrix_world.to_translation() + up_vec
            head_2d = world_to_camera_view(scene, cam, ref_ob.matrix_world.to_translation())
            tail_2d = world_to_camera_view(scene, cam, tail_3d)
            ratio = scene.render.resolution_y / scene.render.resolution_x
            head_2d.y *= ratio
            tail_2d.y *= ratio

        vec_from_corner_2d = (tail_2d - head_2d).normalized()
        up_vec_2d = Vector((0,1))
        # angle = acos(up_vec_2d.dot(vec_from_corner_2d)) ## equivalent but not signed!
        angle = up_vec_2d.angle_signed(vec_from_corner_2d)

        ## Axis camera aim (seem slightly off)
        # rot_axis = Vector((0, 0, -1))
        # rot_axis.rotate(cam.matrix_world)

        ## Axis camera origin -> pivot
        rot_axis = head_world_coord - cam.matrix_world.translation        
        mat = rotate_matrix_around_pivot(mat, angle, head_world_coord, rot_axis)

    else:
        if bone:
            ## Use mid bone to better follow movement
            mat.translation = ref_ob.matrix_world @ ((bone.tail + bone.head) / 2) # Mid bone
        else:
            mat.translation = ref_ob.matrix_world
    
    ## change/adapt scale 
    # mat_scale = Matrix.Scale(10, 4) # maybe move above mesh condition
    # mat = mat @ mat_scale
    obj.matrix_world = mat



def create_plane(name='Plane', collection=None):
    '''Create a plane using pydata
    collection: link in passed collection, else do not link in scene
    '''
    x = 100.0
    y = 100.0
    vert = [(-x, -y, 0.0), (x, -y, 0.0), (-x, y, 0.0), (x, y, 0.0)]
    fac = [(0, 1, 3, 2)]
    pl_data = bpy.data.meshes.new(name)
    pl_data.from_pydata(vert, [], fac)
    pl_obj = bpy.data.objects.new(name, pl_data)
    # collection = bpy.context.collection
    if collection:
        collection.objects.link(pl_obj)
    return pl_obj

def intersect_with_tesselated_plane(point, origin, face_co):
    '''
    face_co: World face coordinates
    '''
    
    tri = None
    for tri_idx in tessellate_polygon([face_co]):
        tri = [face_co[i] for i in tri_idx]
        tri_indices = [i for i in tri_idx]
        hit_location = intersect_line_plane(origin, point, sum((Vector(v) for v in tri), Vector()) / 3, triangle_normal(*tri))  
        if intersect_point_tri(hit_location, *tri):
            break
    
    return np.array(hit_location), tri, tri_indices


def get_bone_head_tail_2d(posebone, scene=None, cam=None) -> tuple[Vector, Vector]:
    '''Get 2D vectors in camera view of bone head and tails
    return tuple of 2d vectors (head_2d and tail_2d) 
    '''
    scene = scene or bpy.context.scene
    cam = cam or scene.camera

    arm = posebone.id_data

    # Get 3D locations of head and tail
    head_3d = arm.matrix_world @ posebone.head
    tail_3d = arm.matrix_world @ posebone.tail

    # Convert 3D locations to 2D
    head_2d = world_to_camera_view(scene, cam, head_3d)
    tail_2d = world_to_camera_view(scene, cam, tail_3d)
    
    ratio = scene.render.resolution_y / scene.render.resolution_x
    head_2d.y *= ratio
    tail_2d.y *= ratio
    
    return Vector((head_2d.x, head_2d.y)), Vector((tail_2d.x, tail_2d.y))


def rotate_matrix_around_pivot(matrix, angle, pivot, axis):
    '''Rotate a given matrix by a CW angle around pivot on a given axis
    matrix (Matrix): the matrix to rotate
    angle (Float, Radians): the angle in radians
    pivot (Vector3): the pivot 3D coordinate
    axis (Vector3): the vector axis of rotation
    '''

    # Convert angle to radians ?
    # angle = math.radians(angle)

    # Create a rotation matrix
    rot_matrix = Matrix.Rotation(angle, 4, axis)

    # Create translation matrices
    translate_to_origin = Matrix.Translation(-pivot)
    translate_back = Matrix.Translation(pivot)

    # Combine the transformations : The order of multiplication is important
    new_matrix = translate_back @ rot_matrix @ translate_to_origin @ matrix

    return new_matrix


# --- GREASE PENCIL

def get_gp_draw_plane(obj=None):
    ''' return tuple with plane coordinate and normal
    of the curent drawing according to geometry'''

    if obj is None:
        obj = bpy.context.object

    settings = bpy.context.scene.tool_settings
    orient = settings.gpencil_sculpt.lock_axis #'VIEW', 'AXIS_Y', 'AXIS_X', 'AXIS_Z', 'CURSOR'
    loc = settings.gpencil_stroke_placement_view3d #'ORIGIN', 'CURSOR', 'SURFACE', 'STROKE'
    
    mat = obj.matrix_world
    plane_no = Vector((0.0, 0.0, 1.0))
    plane_co = mat.to_translation()

    # -> orientation
    if orient == 'VIEW':
        mat = bpy.context.scene.camera.matrix_world 
    # -> placement
    if loc == "CURSOR":
        plane_co = bpy.context.scene.cursor.location
        mat = bpy.context.scene.cursor.matrix

    elif orient == 'AXIS_Y':#front (X-Z)
        plane_no = Vector((0,1,0))

    elif orient == 'AXIS_X':#side (Y-Z)
        plane_no = Vector((1,0,0))

    elif orient == 'AXIS_Z':#top (X-Y)
        plane_no = Vector((0,0,1))

    plane_no.rotate(mat)
    
    return plane_co, plane_no


## --- Animation

def following_keys(forward=True, all_keys=False) -> list:# -> list[int] | list | None:
    '''Return a list of int or an empty list'''
    direction = 1 if forward else -1
    cur_frame = bpy.context.scene.frame_current
    settings = bpy.context.scene.gp_interpo_settings

    scn = bpy.context.scene
    if forward:
        limit = scn.frame_preview_end if scn.use_preview_range else scn.frame_end
    else:
        limit = scn.frame_preview_start if scn.use_preview_range else scn.frame_start

    frames = []
    if settings.mode == 'FRAME':
        jump = settings.padding * direction
        if all_keys:
            limit += direction # offset by one for limit to be in range
            return list(range(cur_frame + jump , limit, jump))

        else:
            return [cur_frame + jump]
    
    elif settings.mode == 'GPKEY':
        layers = bpy.context.object.data.layers
        frames = [f.frame_number for l in layers for f in l.frames]
    
    elif settings.mode == 'RIGKEY':
        col = settings.target_collection
        if not col:
            col = bpy.context.scene.collection            
        
        objs = [o for o in col.all_objects if o.type == 'ARMATURE']
        # Add camera moves detection
        objs += [bpy.context.scene.camera]

        for obj in objs:
            print(obj.name)
            if not obj.animation_data or not obj.animation_data.action:
                continue
            frames += [k.co.x for fc in obj.animation_data.action.fcurves for k in fc.keyframe_points]
        
    
    if not frames:
        return []

    # Sort frames (invert if looking backward)
    frames.sort(reverse=not forward)

    if all_keys:
        frames = list(set(frames))
        if forward:
            frame_list = [int(f) for f in frames if f > cur_frame and f <= limit]
        else:
            frame_list =  [int(f) for f in frames if f < cur_frame and f >= limit]
        return frame_list

    if forward:
        new = next((f for f in frames if f > cur_frame), None)
    else:
        new = next((f for f in frames if f < cur_frame), None)
    if new is None:
        return []
    return [int(new)]


def index_list_from_bools(bool_list) -> list:
    '''Receive a list of boolean
    Return a list of sublists of indices where there is a continuity of True.
    e.g., [True, True, False, True] will return [[0,1][3]]
    '''
    result = []
    current_sublist = []

    for i, value in enumerate(bool_list):
        if value:
            current_sublist.append(i)
        elif current_sublist:
            result.append(current_sublist)
            current_sublist = []

    if current_sublist:
        result.append(current_sublist)

    return result

## -- animation

def is_animated(obj):
    return True