ThirdParty/Rignet_utils/mst_utils.py

#-------------------------------------------------------------------------------
# Name:        mst_utils.py
# Purpose:     utilize functions for skeleton generation
# RigNet Copyright 2020 University of Massachusetts
# RigNet is made available under General Public License Version 3 (GPLv3), or under a Commercial License.
# Please see the LICENSE README.txt file in the main directory for more information and instruction on using and licensing RigNet.
#-------------------------------------------------------------------------------

import sys
import numpy as np
from .rig_parser import TreeNode
from .rig_parser import Skel
import torch

def inside_check(pts, vox):
    """
    Check where points are inside or outside the mesh based on its voxelization.
    :param pts: points to be checked
    :param vox: voxelized mesh
    :return: internal points, and index of them in the input array.
    """
    vc = (pts - vox.translate) / vox.scale * vox.dims[0]
    vc = np.round(vc).astype(int)
    ind1 = np.logical_and(np.all(vc >= 0, axis=1), np.all(vc < 88, axis=1))
    vc = np.clip(vc, 0, 87)
    ind2 = vox.data[vc[:, 0], vc[:, 1], vc[:, 2]]
    ind = np.logical_and(ind1, ind2)
    pts = pts[ind]
    return pts, np.argwhere(ind).squeeze()


def sample_on_bone(p_pos, ch_pos):
    """
    sample points on a bone
    :param p_pos: parent joint position
    :param ch_pos: child joint position
    :return: a array of samples on this bone.
    """
    ray = ch_pos - p_pos
    bone_length = np.sqrt(np.sum((p_pos - ch_pos) ** 2))
    num_step = np.round(bone_length / 0.01)
    i_step = np.arange(1, num_step + 1)
    unit_step = ray / (num_step + 1e-30)
    unit_step = np.repeat(unit_step[np.newaxis, :], num_step, axis=0)
    res = p_pos + unit_step * i_step[:, np.newaxis]
    return res


def minKey(key, mstSet, nV):
    # Initilaize min value
    min = sys.maxsize
    for v in range(nV):
        if key[v] < min and mstSet[v] == False:
            min = key[v]
            min_index = v
    return min_index 
    
def primMST_normal(graph, init_id, normal_matrix):
    """
    Modified Prim's algorithm to generate a minimum spanning tree (MST).
    :param graph: pairwise cost matrix
    :param init_id: init node ID as root
    :return: parent array, key array, init_id
    """
    nV = graph.shape[0]
    key = [sys.maxsize] * nV
    parent = [None] * nV
    mstSet = [False] * nV
    key[init_id] = 0
    parent[init_id] = -1
    previous_normal = np.zeros((nV, 3))

    while not all(mstSet):
        u = minKey(key, mstSet, nV)
        mstSet[u] = True
        if parent[u] >= 0:
            previous_normal[u] = normal_matrix[u, parent[u]]
            updated_normal = np.dot(previous_normal[u], normal_matrix[u, :].T) #1*n
            updated_normal[updated_normal<0]=0
        # print('updated_normal',updated_normal.shape)
            graph[u, :] = graph[u, :] +(1e8*updated_normal**2+1)
            graph[:, u] = graph[:, u] +(1e8*updated_normal**2+1)

        for v in range(nV):
            
            if graph[u, v] > 0 and mstSet[v] is False and key[v] > graph[u, v]:
                key[v] = graph[u, v]
                parent[v] = u


    return parent, key, init_id   


def loadSkel_recur(p_node, parent_id, joint_name, joint_pos, parent):
    """
    Converst prim algorithm result to our skel/info format recursively
    :param p_node: Root node
    :param parent_id: parent name of current step of recursion.
    :param joint_name: list of joint names
    :param joint_pos: joint positions
    :param parent: parent index returned by prim alg.
    :return: p_node (root) will be expanded to linked with all joints
    """
    for i in range(len(parent)):
        if parent[i] == parent_id:
            if joint_name is not None:
                ch_node = TreeNode(joint_name[i], tuple(joint_pos[i]))
            else:
                ch_node = TreeNode('joint_{}'.format(i), tuple(joint_pos[i]))
            p_node.children.append(ch_node)
            ch_node.parent = p_node
            loadSkel_recur(ch_node, i, joint_name, joint_pos, parent)


def unique_rows(a):
    """
    remove repeat rows from a numpy array
    """
    a = np.ascontiguousarray(a)
    unique_a = np.unique(a.view([('', a.dtype)]*a.shape[1]))
    return unique_a.view(a.dtype).reshape((unique_a.shape[0], a.shape[1]))


def increase_cost_for_outside_bone(cost_matrix, joint_pos, vox):
    """
    increase connectivity cost for bones outside the meshs
    """
    for i in range(len(joint_pos)):
        for j in range(i+1, len(joint_pos)):
            bone_samples = sample_on_bone(joint_pos[i], joint_pos[j])
            bone_samples_vox = (bone_samples - vox.translate) / vox.scale * vox.dims[0]
            bone_samples_vox = np.round(bone_samples_vox).astype(int)

            ind1 = np.logical_and(np.all(bone_samples_vox >= 0, axis=1), np.all(bone_samples_vox < vox.dims[0], axis=1))
            bone_samples_vox = np.clip(bone_samples_vox, 0, vox.dims[0]-1)
            ind2 = vox.data[bone_samples_vox[:, 0], bone_samples_vox[:, 1], bone_samples_vox[:, 2]]
            in_flags = np.logical_and(ind1, ind2)
            outside_bone_sample = np.sum(in_flags == False)

            if outside_bone_sample > 1:
                cost_matrix[i, j] = 2 * outside_bone_sample
                cost_matrix[j, i] = 2 * outside_bone_sample
            if np.abs(joint_pos[i, 0]) < 2e-2 and np.abs(joint_pos[j, 0]) < 2e-2:
                cost_matrix[i, j] *= 0.5
                cost_matrix[j, i] *= 0.5
    return cost_matrix

def increase_cost_for_outside_bone_tensor(cost_matrix, joint_pos, vox,resolution=64):
    """
    increase connectivity cost for bones outside the meshs
    vox is a tensor with size(N,3), N is the number of voxels that inside the mesh, range (0,64)
    """
    
    vox = torch.clamp(vox, 0, resolution-1).long()
    for i in range(len(joint_pos)):
        for j in range(i+1, len(joint_pos)):
            bone_samples = sample_on_bone(joint_pos[i], joint_pos[j]) # return coordinates of points on the bone
            bone_samples_vox = bone_samples * (resolution/2) + (resolution/2)
            bone_samples_vox = np.round(bone_samples_vox).astype(int)
            bone_samples_vox = np.clip(bone_samples_vox, 0, resolution-1)
            
            vox_remap = torch.zeros((resolution,resolution,resolution))
            vox_remap[vox[:,0],vox[:,1],vox[:,2]] = 1
            vox_remap = vox_remap.numpy()
            inside_index = vox_remap[bone_samples_vox[:,0],bone_samples_vox[:,1],bone_samples_vox[:,2]]
            outside_bone_sample = np.sum(inside_index == 0)

            
            # check the intersection of the bone with the mesh

            if outside_bone_sample > 1:
                cost_matrix[i, j] = 2 * outside_bone_sample
                cost_matrix[j, i] = 2 * outside_bone_sample
            if np.abs(joint_pos[i, 0]) < 2e-2 and np.abs(joint_pos[j, 0]) < 2e-2:
                cost_matrix[i, j] *= 0.5
                cost_matrix[j, i] *= 0.5
    return cost_matrix