get_loops.py

from typing import Optional, Sequence, Union, List
import numpy as np
from scipy.interpolate import interp1d
import matplotlib.path as pltpath  # for the area-sampling function
import polyscope as ps

import utils
from geometry_utils import unit, extract_useful_data_from_plane, polygon_area
import meshlab_poisson_reco # has thlog
import slice_with_pyvista # has the slicer
from thlog import *

# status codes for get_face_intersections_with_plane
PLANE_CUTS_FACE_AT_TWO_SIDES = 0
PLANE_CUTS_FACE_AT_ONE_SIDE_AND_TOUCHES_ONE_V = 1
PLANE_TOUCHES_FACE_AT_ONE_V = 2
PLANE_TOUCHES_FACE_AT_TWO_V = 3
PLANE_CONTAINS_WHOLE_FACE = 4

# logger for the get_loops module
thlog = Thlogger(LOG_INFO, VIZ_DEBUG, "slice", imports=[meshlab_poisson_reco.thlog])


class MeshOneSlice:
    """ 
    Caches some info about a single set of polylines created by
    slicing a mesh with a SINGLE plane.

    Initialization parameters:
    - plane: the slice plane this slice is associated with. Represented as an
        array of shape (4, 3) where plane_n = plane[0], plane_p = plane[1], and
        plane[2], plane[3] are on-plane coordinate system x axis and y axis
        vectors. (plane_n is the plane's normal and plane_p is a point on the
        plane)
    - mesh_fname: Pass here an .obj file path to slice for ground truth loops.
        To NOT slice any mesh and simply populate values based on a
        predicted set of loops associated with the plane, set this to None.
        (which is the use case for when we receive loops generated by a
        model that predicts loops)
    - predicted_loops: This should be a tuple consisting of
        (loop points, loop connectivity array, loop normals (this may be None))
    """
    def __init__(self, plane: np.ndarray, mesh_fname=None, predicted_loops=None):
        self.from_slicing_real_mesh = False
        self.plane = plane
        assert plane.shape == (4, 3), \
            "plane array must have shape (4,3), consisting of plane normal, plane point, on-plane coordinate system x axis and y axis vectors "
        assert np.isclose(
            np.abs(np.dot(
                unit(np.cross(plane[2], plane[3])), unit(plane[0]))), 1), \
                    "supplied on-plane coordinate system doesn't work!"
        # ^^ plane_n, plane_p, onplane_xaxis, onplane_yaxis (i.e. the x and y
        # axis basis vectors for a 2D coordinate system we define on the plane.
        # Here we check if the x and y vectors' cross product will line up with
        # the normal (i.e. cosine sim will either be 1 or -1 ))

        if mesh_fname is not None:
            self.from_slicing_real_mesh = True
            loop_pts, loop_conn_arr, normals_at_loop_pts = \
                slice_with_pyvista.slice_mesh_using_pyvista(mesh_fname,
                    plane)
        elif predicted_loops is not None:
            loop_pts, loop_conn_arr, normals_at_loop_pts = predicted_loops
        else:
            raise ValueError("At least one of hemesh or predicted_loops should be non-None")
        self.loop_pts: np.ndarray = loop_pts
        self.loop_conn_arr: np.ndarray = loop_conn_arr
        # for caching later when other functions are run (such as the normal
        # estimator)
        self.disjoint_loop_conns: Optional[Sequence[np.ndarray]] = None
        self.disjoint_loop_lengths: Optional[Sequence[float]] = None
        self.loop_segment_normals: Optional[np.ndarray] = normals_at_loop_pts
        self.loop_segment_vectors: Optional[np.ndarray] = None
        self.sign_flip_per_loop: Optional[np.ndarray] = None  # for inside-outside information
        # fill in the rest of the arrays
        distinguish_disjoint_loops(self)
    
    def show_in_polyscope(self):
        if thlog.guard(VIZ_INFO, needs_polyscope=True):
            ps.register_point_cloud("slice debug", self.loop_pts)
            ps.show()
    


def distinguish_disjoint_loops(mesh_slice: MeshOneSlice):
    """
    Splits up the mesh_slice result into disjoint loops, since a slice plane
    through a mesh may result in multiple closed loops being formed (i.e. if it
    slices through the fingers of a hand mesh). Populates the
    disjoint_loop_conns member of the MeshOneSlice object. Since we scan through
    all segments, we might as well populate the segment vectors of the loop's
    segments in loop_segment_vectors.

    Reads the following fields in mesh_slice:
    - loop_pts, loop_conn_arr
    Fills in the following fields in mesh_slice:
    - disjoint_loop_conns, loop_segment_vectors
    """

    entering_new_loop = True
    curr_polygonstart = -1

    loop_pts = mesh_slice.loop_pts
    loop_conn_arr = mesh_slice.loop_conn_arr

    loop_seg_vectors = np.zeros_like(loop_pts)
    curr_loop_conn = []
    disjoint_loop_conns = []
    
    for segment in loop_conn_arr:
        curr_i = segment[0]
        next_i = segment[1]
        if entering_new_loop:
            curr_polygonstart = curr_i
            entering_new_loop = False
        
        curr_loop_conn.append(segment)
        seg_vec = loop_pts[next_i] - loop_pts[curr_i]
        loop_seg_vectors[curr_i] = seg_vec

        if next_i == curr_polygonstart:
            # we've wrapped around and finished a polygon, so skip this iter
            # and remember to update the curr_polygonstart index next iter

            # add in the last segment for the conn of the current loop
            # curr_loop_conn.append(segment)
            # add the populated current loop to the list of disjoint loops
            disjoint_loop_conns.append(curr_loop_conn)
            # and reset curr_ for the next loop
            curr_loop_conn = []
            entering_new_loop = True
            continue
    mesh_slice.disjoint_loop_conns = list(map(np.array, disjoint_loop_conns))
    mesh_slice.loop_segment_vectors = loop_seg_vectors
    return mesh_slice


def sort_loops_by_decreasing_area(mesh_slice: MeshOneSlice, change_to_plane_basis, plane_p_in_plane_basis):
    """ sort disjoint loops in a MeshOneSlice in order of descending area
    (so that the biggest loop comes first in the sequence, and so on)
    Modifies the mesh_slice in place but also returns back (a reference to) it.
    """
    loop_pts_in_plane_basis = \
        np.transpose(np.dot(change_to_plane_basis, np.transpose(mesh_slice.loop_pts))) - plane_p_in_plane_basis

    assert mesh_slice.disjoint_loop_conns is not None
    disjoint_loop_conns = mesh_slice.disjoint_loop_conns
    def __calculate_area(loop_conn_arr) -> float:
        loop_pts_in_plane_basis_this_loop = loop_pts_in_plane_basis[loop_conn_arr[:, 0]]
        return polygon_area(loop_pts_in_plane_basis_this_loop[:, 0], loop_pts_in_plane_basis_this_loop[:, 1])

    sorted_indices = sorted(list(range((len(disjoint_loop_conns)))), 
        key=lambda i: __calculate_area(disjoint_loop_conns[i]), reverse=True)

    mesh_slice.disjoint_loop_conns = [mesh_slice.disjoint_loop_conns[i] for i in sorted_indices]
    # mesh_slice.disjoint_loop_lengths = [mesh_slice.disjoint_loop_lengths[i] for i in sorted_indices]
    return mesh_slice, loop_pts_in_plane_basis


def calculate_inside_outside_loops(mesh_slice: MeshOneSlice):
    """
    populates the field mesh_slice.sign_flip_per_loop, an int array where each
    value is either 1 or -1 (related to whether the winding number is even or
    odd, i.e. a loop B completely contained in a loop A gets the inverse of loop
    A's sign; the outermost loop gets 1)
    """
    disjoint_loop_conns = mesh_slice.disjoint_loop_conns
    assert disjoint_loop_conns is not None

    n_disjoint_loops = len(disjoint_loop_conns)
    if n_disjoint_loops > 1:
        # inside-outside check; potentially expensive
        _, _, _, _, _, change_to_plane_basis, plane_p_in_plane_basis = extract_useful_data_from_plane(mesh_slice.plane)
        mesh_slice, loop_pts_in_plane_basis = sort_loops_by_decreasing_area(mesh_slice, change_to_plane_basis, plane_p_in_plane_basis)
        disjoint_loop_conns = mesh_slice.disjoint_loop_conns
        assert disjoint_loop_conns is not None

        # ^ this loop_pts_in_plane_basis has (N, 3), we chop off the last coord
        # which should be zero
        loop_pts_in_2d = loop_pts_in_plane_basis[:, :2]
        # then for each loop in this sorted mesh slice, check all smaller loops
        sign_flip_per_loop = np.ones(n_disjoint_loops, dtype=int)
        loops_as_paths = [pltpath.Path(loop_pts_in_2d[loop_conn[:, 0]], closed=True) for loop_conn in disjoint_loop_conns]
        
        for path_i, path in enumerate(loops_as_paths):
            for smaller_path_i in range(path_i+1, n_disjoint_loops):
                if path.contains_path(loops_as_paths[smaller_path_i]):
                    sign_flip_per_loop[smaller_path_i] = -sign_flip_per_loop[smaller_path_i]

        mesh_slice.sign_flip_per_loop = sign_flip_per_loop
    else:
        mesh_slice.sign_flip_per_loop = np.array([1], dtype=int)
    return mesh_slice


def estimate_loop_segment_normals_simple(mesh_slice: MeshOneSlice):
    """
    A normal estimation routine. Estimates loop normals based on the slice plane
    normal and the loop segment vectors (i.e. takes cross product). Because of
    this estimation method, all the estimated loop segment normals will end up
    on planes parallel to the original slice plane.
    Assumes that distinguish_disjoint_loops has already been run.

    Parameters:
    - mesh_slice: MeshOneSlice object containing the cached data of the sliced
        mesh, the slice plane, and the polyline(s) resulting from this slice.
    
    Returns: 
    - mesh_slice: updated mesh_slice object (though this also updates
        the mesh_slice object in-place)
    
    Reads the following fields in mesh_slice:
    - plane, loop_pts, loop_conn_arr, loop_segment_vectors, disjoint_loop_conns
    Changes the following fields in mesh_slice:
    - loop_segment_normals
    """
    mesh_slice = calculate_inside_outside_loops(mesh_slice)
    plane = mesh_slice.plane

    loop_pts = mesh_slice.loop_pts
    loop_conn_arr = mesh_slice.loop_conn_arr

    loop_seg_normals = np.zeros_like(loop_pts)
    loop_seg_vectors = mesh_slice.loop_segment_vectors
    assert loop_seg_vectors is not None

    plane_n = plane[0]
    plane_p = plane[1]
    
    for segment in loop_conn_arr:
        curr_i = segment[0]
        next_i = segment[1]
        
        seg_vec = loop_seg_vectors[curr_i]
        seg_norm = np.cross(plane_n, seg_vec)
        loop_seg_normals[curr_i] = seg_norm
        
    # assert len(polygon_starts) == len(polygon_ends)
    # find a single comparison point that is on the "same side" of all possible
    # loops obtainable from this mesh (to determine the orientation differences
    # of the loops)
    # mean_mesh_pt = np.mean(hemesh.vertices, axis=0)
    
    ref_point_offset = 200. # this is probably big enough for all meshes we use
    ref_point = plane_p + ref_point_offset * (plane_n / np.linalg.norm(plane_n))
    
    # then sum up the cross-prods of each vector segment with the reference
    # point, and check the sign against the plane normal. Flip the sign of the
    # loops to make that comparison (with the plane normal) be the same sign for
    # all disjoint loops.
    assert mesh_slice.disjoint_loop_conns is not None
    assert mesh_slice.sign_flip_per_loop is not None # for inside-outside info
    disjoint_loop_conns = mesh_slice.disjoint_loop_conns
    for loop_conn, sign_flip in zip(disjoint_loop_conns, mesh_slice.sign_flip_per_loop):
        sum_crosses = np.zeros(3)
        for segment in loop_conn:
            curr_i = segment[0]
            next_i = segment[1]
            sum_crosses += np.cross((loop_pts[curr_i] - ref_point), 
                                    (loop_pts[next_i] - ref_point))
        similarity_to_plane_n = np.dot(plane_n, sum_crosses)
        if similarity_to_plane_n >= 0:
            # flip the sign of all the normals of this loop
            for segment in loop_conn:
                loop_seg_normals[segment[0]] *= (-1)
        
        # finally, flip sign according to inside-outside calculation
        for segment in loop_conn:
            loop_seg_normals[segment[0]] *= sign_flip

    mesh_slice.loop_segment_normals = loop_seg_normals
    return mesh_slice


def sample_from_loops_and_normals(mesh_slice: MeshOneSlice, n_points: int,
        separate_by_disjoint_loops=False, 
        distribute_points_by_loop_length=True, 
        min_n_points_per_loop=-1,
        normal_lerping=True
        ):
    # using the approach from loopcnn-code/loop_util.py; For each disjoint loop,
    # create a scipy.interpolate.interp1d function that takes in a "progress"
    # value, (progress as in distance traveled along the closed polyline from
    # the first point) and returns an interpolated position on the loop. Do the
    # same for the normals.
    """ 
    From loops, creates a pointcloud w/ normals per point by sampling points
    along the loop segments and assigning to them the per-segment normals saved
    in the mesh_slice object from some previously-run normal-extraction routine
    (which could either be normal estimation from slice plane normal and loop
    segment vectors, or ground-truth mesh edge normals matched to the
    appropriate loop segment)
    Assumes that distinguish_disjoint_loops has already been run.

    Parameters:
    - mesh_slice: MeshOneSlice object containing the cached data of the sliced
        mesh, the slice plane, and the polyline(s) resulting from this slice
    - n_points: target number of points to sample
    - (optional, default=False) separate_by_disjoint_loops:
        Specify True in order to obtain lists of arrays of sampled points and
        sampled normals separated by disjoint loops
    - (optional, default=True) distribute_points_by_loop_length:
        Specify False in order to force all disjoint loops to have the 
        same n_points sampled from them. If True, then n_points will be divided/
        be distributed among all disjoint loops, proportional to their length.
        (new 2022/07/16)
    - (optional, default=-1) min_n_points_per_loop:
        When "distribute_points_by_loop_length" is True, specify a 
        nonnegative number here in order to impose a minimum number of points
        representing each loop. This is useful for sampling very short loops
        relative to the rest of the loop lengths in this mesh slice, doing so
        in a way that doesn't cause it to get only 2 or 3 points which would
        be too few points for a good reconstruction/viewing experience.
        Specify a negative number to disable this limit.
    - (optional, default=True) normal_lerping:
        If True, the estimated normals on the sampled points are interpolated
        between the normals of the endpoints of the edge that the sampled
        points come from. (i.e. smoother normals, but less sharp corners.)
        If False, the nearest endpoint normal is used.

    Returns:  tuple (all_sampled_points, all_sampled_normals) where
    - all_sampled_points: np.array (hopefully n_points, 3)
    - all_sampled_normals:  np.array (hopefully n_points, 3), one normal per pt

    Reads the following fields in mesh_slice:
    - disjoint_loop_conns, loop_segment_vectors, loop_segment_normals, loop_pts
    (the loop segment normals could either have been taken from ground truth
    normals, or they could have been replaced by fake calculated normals
    from estimate_loop_segment_normals_simple.)
    """

    disjoint_loop_conns = mesh_slice.disjoint_loop_conns
    loop_seg_vectors = mesh_slice.loop_segment_vectors
    loop_seg_normals = mesh_slice.loop_segment_normals
    loop_pts = mesh_slice.loop_pts

    if (disjoint_loop_conns is None) or \
        (loop_seg_vectors is None) or \
        (loop_seg_normals is None) or \
        (loop_pts is None):
        raise ValueError("Must run distinguish_disjoint_loops first!")

    loop_seg_lengths = np.linalg.norm(loop_seg_vectors, axis=1)
    calc_loop_length = lambda closed_loop_conn_arr: \
        np.sum([loop_seg_lengths[seg[0]] for seg in closed_loop_conn_arr])
    disjoint_loop_lengths = np.array(list(map(calc_loop_length, disjoint_loop_conns)))
    proportion_of_total_polyline_length = \
        disjoint_loop_lengths / np.sum(disjoint_loop_lengths)
    
    if separate_by_disjoint_loops:
        all_sampled_points: Union[List[np.ndarray], Optional[np.ndarray]] = []
        all_sampled_normals: Union[List[np.ndarray], Optional[np.ndarray]] = []
    else:
        all_sampled_points: Union[List[np.ndarray], Optional[np.ndarray]] = None
        all_sampled_normals: Union[List[np.ndarray], Optional[np.ndarray]] = None

    for loop_conn_arr, loop_length, proportion in zip(
        disjoint_loop_conns, 
        disjoint_loop_lengths, 
        proportion_of_total_polyline_length):
        
        # newly added optional arg (2022/07/16)
        if distribute_points_by_loop_length:
            # distribute proportional to length (default behavior)
            n_points_this_loop = round(proportion * n_points)
            if min_n_points_per_loop:
                n_points_this_loop = max(min_n_points_per_loop, n_points_this_loop)
        else:
            # force n_points always, don't distribute by length
            n_points_this_loop = n_points
        
        first_i_of_segs_no_backtostart = loop_conn_arr[:, 0]
        first_i_of_segs = np.append(
            first_i_of_segs_no_backtostart, first_i_of_segs_no_backtostart[0])
        
        lerpfn_in = np.cumsum(loop_seg_lengths[first_i_of_segs_no_backtostart])
        lerpfn_in[-1] = loop_length # apparently avoids rounding errors
        lerpfn_in = np.insert(lerpfn_in, 0, 0)
        lerpfn_progress = np.linspace(0, loop_length, n_points_this_loop)

        # lerp func for point-sampling
        lerpfn_out_px = loop_pts[first_i_of_segs, 0]
        lerpfn_out_py = loop_pts[first_i_of_segs, 1]
        lerpfn_out_pz = loop_pts[first_i_of_segs, 2]
        lerpfn_px = interp1d(lerpfn_in, lerpfn_out_px)
        lerpfn_py = interp1d(lerpfn_in, lerpfn_out_py)
        lerpfn_pz = interp1d(lerpfn_in, lerpfn_out_pz)

        # lerp func for normal sampling (either step function, which would grant
        # the same normal (associated with each segment) to ALL points sampled
        # from the same segment, or lin interp, which would smoothly lerp the
        # normal from end to end.)

        lerpfn_out_nx = loop_seg_normals[first_i_of_segs, 0]
        lerpfn_out_ny = loop_seg_normals[first_i_of_segs, 1]
        lerpfn_out_nz = loop_seg_normals[first_i_of_segs, 2]
        # nearest = floor-step func; linear = lerp
        lerp_strategy = 'linear' if normal_lerping else 'nearest'
        lerpfn_nx = interp1d(lerpfn_in, lerpfn_out_nx, kind=lerp_strategy) 
        lerpfn_ny = interp1d(lerpfn_in, lerpfn_out_ny, kind=lerp_strategy)
        lerpfn_nz = interp1d(lerpfn_in, lerpfn_out_nz, kind=lerp_strategy)

        sampled_points_this_loop = np.vstack((
            lerpfn_px(lerpfn_progress),
            lerpfn_py(lerpfn_progress),
            lerpfn_pz(lerpfn_progress)
        ))

        sampled_normals_this_loop = np.vstack((
            lerpfn_nx(lerpfn_progress),
            lerpfn_ny(lerpfn_progress),
            lerpfn_nz(lerpfn_progress)
        ))
        if separate_by_disjoint_loops and isinstance(all_sampled_points, list) and isinstance(all_sampled_normals, list):
            all_sampled_points.append(sampled_points_this_loop)
            all_sampled_normals.append(sampled_normals_this_loop)
        else:
            if all_sampled_points is None or all_sampled_normals is None:
                all_sampled_points = sampled_points_this_loop
                all_sampled_normals = sampled_normals_this_loop
            else:
                all_sampled_points = np.hstack(
                    (all_sampled_points, sampled_points_this_loop))
                all_sampled_normals = np.hstack(
                    (all_sampled_normals, sampled_normals_this_loop))
    if separate_by_disjoint_loops and isinstance(all_sampled_points, list) and isinstance(all_sampled_normals, list):
        return list(map(np.transpose, all_sampled_points)), \
               list(map(np.transpose, all_sampled_normals))
    else:
        assert isinstance(all_sampled_points, np.ndarray) and isinstance(all_sampled_normals, np.ndarray)
        return np.transpose(all_sampled_points), \
               np.transpose(all_sampled_normals)

def sample_points_to_fill_polygons(mesh_slice: MeshOneSlice, n_bbox_points: int, flip_normals: bool):
    """
    Sample points in the area of each closed polygon in the given slice, and
    give those points the normal of the plane (or the flipped normal). Useful
    for sampling points to form a "cap" to prevent holes at the top and bottom
    in reconstruction results.
    Parameters:
    - mesh_slice: slice object, with disjoint_loop_conns populated
    - n_bbox_points: number of points to sample within the bounding box of each
        polygon to then filter out for points actually inside the polygon
    - flip_normals: bool; whether to take the plane normal or its negative
    Returns: (extra sampled points, their normals) 
    (it's the caller's responsibility to then concatenate this with the rest of
    the points and normals for viz.)
    """

    plane_n, plane_p, _, _, plane_basis, change_to_plane_basis, plane_p_in_plane_basis = \
        extract_useful_data_from_plane(mesh_slice.plane)
    loop_pts_in_plane_basis = \
        np.transpose(np.dot(change_to_plane_basis, np.transpose(mesh_slice.loop_pts))) -  plane_p_in_plane_basis
    loop_pts_in_2d = loop_pts_in_plane_basis[:, :2]

    # the loop conn arr of each closed polygon is guaranteed to be in order of
    # vertices around the polygon (and always clockwise), so we can do this
    # (it's guaranteed because of the way distinguish_disjoint_loops is run)
    assert mesh_slice.disjoint_loop_conns is not None
    assert mesh_slice.sign_flip_per_loop is not None  # for filtering points inside/outside polygons
    polygons = [loop_pts_in_2d[disjoint_loop_conn[:, 0]] for disjoint_loop_conn in mesh_slice.disjoint_loop_conns]
    
    # find the bounding box for each polygon. (top left corner, bottom right corner)
    bboxes = [(np.array([np.min(pts[:,0]), np.min(pts[:,1])])
              ,np.array([np.max(pts[:,0]), np.max(pts[:,1])])) for pts in polygons]
    
    # sample points inside each bbox
    bbox_sampled_points = [np.reshape(
        np.stack(np.meshgrid(
            np.linspace(pmin[0], pmax[0], n_bbox_points), np.linspace(pmin[1], pmax[1], n_bbox_points))
            , axis=-1)
        , (-1, 2))
        for (pmin, pmax) in bboxes]
    
    # then for each polygon+bbox sample pts, filter these points to only those
    # in the polygon.
    # account for inside-outside (i.e. only count polygons with a sign of 1)    
    sampled_points_inside_polygons = [
        bbox_pts[pltpath.Path(polygon).contains_points(bbox_pts)]
        for polygon, bbox_pts, sign_flip in zip(polygons, bbox_sampled_points, mesh_slice.sign_flip_per_loop) if sign_flip > 0]
    
    # recast these points from the 2D plane basis to 3D world coords
    sampled_points_inside_polygons = [np.transpose(
        np.dot(plane_basis, np.transpose(np.hstack((pts, np.zeros((len(pts), 1))))))) + plane_p 
        for pts in sampled_points_inside_polygons ]

    # then just repeat the plane normal as many times as there are points
    normals_each_polygon = [np.tile(plane_n * ((-1) if flip_normals else 1), (len(pts), 1))
        for pts in sampled_points_inside_polygons
    ]
    
    # merge the point arrays and normal arrays of all polygons into one array 
    return np.concatenate(sampled_points_inside_polygons),\
        np.concatenate(normals_each_polygon)