src/data/asset.py

from collections import defaultdict
from dataclasses import dataclass
import numpy as np
from numpy import ndarray

from typing import Dict, Union, List, Tuple

from .order import Order
from .raw_data import RawData
from .exporter import Exporter

from ..tokenizer.spec import TokenizeInput
from .utils import linear_blend_skinning

import trimesh


@dataclass
class Asset(Exporter):
    '''
    Dataclass to handle data parsed from raw data.
    '''
    
    # data class
    cls: str
    
    # where is this asset from
    path: str
    
    # data file name
    data_name: str
    
    # vertices of the mesh, shape (N, 3), float32
    vertices: ndarray
    
    # normals of vertices, shape (N, 3), float32
    vertex_normals: ndarray
    
    # faces of mesh, shape (F, 3), face id starts from 0 to F-1, int64
    faces: ndarray
    
    # face normal of mesh, shape (F, 3), float32
    face_normals: ndarray
    
    # joints of bones, shape (J, 3), float32
    joints: Union[ndarray, None]=None
    
    # tails of joints, shape (J, 3), float32
    tails: Union[ndarray, None]=None
    
    # skinning of joints, shape (N, J), float32
    skin: Union[ndarray, None]=None
    
    # whether the joint has skin, bool
    no_skin: Union[ndarray, None]=None
    
    # vertex groups
    vertex_groups: Union[Dict[str, ndarray], None]=None
    
    # parents of joints, None represents no parent(a root joint)
    # make sure parent[k] < k
    parents: Union[List[Union[int, None]], None]=None
    
    # names of joints
    names: Union[List[str], None]=None
    
    # sampled vertices, shape (N, 3)
    sampled_vertices: Union[ndarray, None]=None
    
    # sampled normals, shape (N, 3)
    sampled_normals: Union[ndarray, None]=None
    
    # sampled vertex groups, every vertex group should be (N, J)
    sampled_vertex_groups: Union[Dict[str, ndarray], None]=None
    
    # {id: part}, part==None -> a spring token
    parts_bias: Union[Dict[int, Union[str, None]], None]=None
    
    # local coordinate, shape (J, 4, 4)
    matrix_local: Union[ndarray, None]=None
    
    # pose matrix for skinning loss calculation, shape (J, 4, 4)
    pose_matrix: Union[ndarray, None]=None
    
    meta: Union[Dict[str, ...], None]=None
    
    @property
    def N(self):
        '''
        number of vertices
        '''
        return self.vertices.shape[0]
    
    @property
    def F(self):
        '''
        number of faces
        '''
        return self.faces.shape[0]
    
    @property
    def J(self):
        '''
        number of joints
        '''
        return self.joints.shape[0]
    
    def get_matrix(self, matrix_basis: ndarray, matrix_local: Union[ndarray, None]=None):
        '''
        get matrix
        
        matrix_basis: (J, 4, 4)
        '''
        if matrix_local is None:
            assert self.joints is not None
            matrix_local = self.matrix_local
            if matrix_local is None:
                matrix_local = np.zeros((self.J, 4, 4))
                matrix_local[:, 0, 0] = 1.
                matrix_local[:, 1, 1] = 1.
                matrix_local[:, 2, 2] = 1.
                matrix_local[:, 3, 3] = 1.
                for i in range(self.J):
                    matrix_local[i, :3, 3] = self.joints[i]
        
        matrix = np.zeros((self.J, 4, 4))
        for i in range(self.J):
            if i==0:
                matrix[i] = matrix_local[i] @ matrix_basis[i]
            else:
                pid = self.parents[i]
                matrix_parent = matrix[pid]
                matrix_local_parent = matrix_local[pid]
                
                matrix[i] = (
                    matrix_parent @
                    (np.linalg.inv(matrix_local_parent) @ matrix_local[i]) @
                    matrix_basis[i]
                )
        return matrix
    
    def apply_matrix_basis(self, matrix_basis: ndarray):
        '''
        apply a pose to armature
        
        matrix_basis: (J, 4, 4)
        '''
        matrix_local = self.matrix_local
        if matrix_local is None:
            matrix_local = np.zeros((self.J, 4, 4))
            matrix_local[:, 0, 0] = 1.
            matrix_local[:, 1, 1] = 1.
            matrix_local[:, 2, 2] = 1.
            matrix_local[:, 3, 3] = 1.
            for i in range(self.J):
                matrix_local[i, :3, 3] = self.joints[i].copy()
        
        matrix = self.get_matrix(matrix_basis=matrix_basis, matrix_local=matrix_local)
        self.joints = matrix[:, :3, 3].copy()
        vertices = linear_blend_skinning(self.vertices, matrix_local, matrix, self.skin, pad=1, value=1.)
        # update matrix_local
        self.matrix_local = matrix.copy()

        # change tails
        if self.tails is not None:
            t_skin = np.eye(self.J)
            self.tails = linear_blend_skinning(self.tails, matrix_local, matrix, t_skin, pad=1, value=1.)
        # in accordance with trimesh's normals
        mesh = trimesh.Trimesh(vertices=vertices, faces=self.faces, process=False)
        self.vertices = vertices
        self.vertex_normals = mesh.vertex_normals.copy()
        self.face_normals = mesh.face_normals.copy()
    
    def set_order_by_names(self, new_names: List[str]):
        assert len(new_names) == len(self.names)
        name_to_id = {name: id for (id, name) in enumerate(self.names)}
        new_name_to_id = {name: id for (id, name) in enumerate(new_names)}
        perm = []
        new_parents = []
        for (new_id, name) in enumerate(new_names):
            perm.append(name_to_id[name])
            pid = self.parents[name_to_id[name]]
            if new_id == 0:
                assert pid is None, 'first bone is not root bone'
            else:
                pname = self.names[pid]
                pid = new_name_to_id[pname]
                assert pid < new_id, 'new order does not form a tree'
            new_parents.append(pid)
        
        if self.joints is not None:
            self.joints = self.joints[perm]
        self.parents = new_parents
        if self.tails is not None:
            self.tails = self.tails[perm]
        if self.skin is not None:
            self.skin = self.skin[:, perm]
        if self.no_skin is not None:
            self.no_skin = self.no_skin[perm]
        if self.matrix_local is not None:
            self.matrix_local = self.matrix_local[perm]
        self.names = new_names
    
    def set_order(self, order: Order):
        if self.names is None or self.parents is None:
            return
        new_names, self.parts_bias = order.arrange_names(cls=self.cls, names=self.names, parents=self.parents)
        self.set_order_by_names(new_names=new_names)
    
    def collapse(self, keep: List[str]):
        dsu = [i for i in range(self.J)]
        
        def find(x: int) -> int:
            if dsu[x] == x:
                return x
            y = find(dsu[x])
            dsu[x] = y
            return y
        
        def merge(x: int, y: int):
            dsu[find(x)] = find(y)

        if self.tails is not None:
            new_tails = self.tails.copy()
        else:
            new_tails = None
        if self.skin is not None:
            new_skin = self.skin.copy()
        else:
            new_skin = None

        if self.no_skin is not None:
            new_no_skin = self.no_skin.copy()
        else:
            new_no_skin = None
        
        if self.matrix_local is not None:
            matrix_local = self.matrix_local.copy()
        else:
            matrix_local = None
        new_names = []
        new_parents = []
        perm = []
        new_name_to_id = {}
        tot = 0
        for (i, name) in enumerate(self.names):
            if name in keep:
                new_names.append(name)
                new_name_to_id[name] = tot
                tot += 1
                perm.append(i)
                pid = self.parents[i]
                if pid is None:
                    new_parents.append(None)
                else:
                    pid = find(pid)
                    new_parents.append(new_name_to_id[self.names[pid]])
                continue
            assert i != 0, 'cannot remove root'
            id = find(i)
            pid = find(self.parents[id])
            # be careful !
            # do not copy tail here because you dont know which child to inherit from
            if new_skin is not None:
                new_skin[:, pid] += new_skin[:, id]
            if new_no_skin is not None:
                new_no_skin[pid] &= new_no_skin[id]
            merge(id, pid)
        
        if new_tails is not None:
            new_tails = new_tails[perm]
        if new_skin is not None:
            new_skin = new_skin[:, perm]
        if new_no_skin is not None:
            new_no_skin = new_no_skin[perm]
        if matrix_local is not None:
            matrix_local = matrix_local[perm]
        
        if self.joints is not None:
            self.joints = self.joints[perm]
        self.parents        = new_parents
        self.tails          = new_tails
        self.skin           = new_skin
        self.no_skin        = new_no_skin
        self.names          = new_names
        self.matrix_local   = matrix_local
    
    @staticmethod
    def from_raw_data(
        raw_data: RawData,
        cls: str,
        path: str,
        data_name: str,
    ) -> 'Asset':
        '''
        Return an asset initialized from raw data and do transform.
        '''
        return Asset(
            cls=cls,
            path=path,
            data_name=data_name,
            vertices=raw_data.vertices,
            vertex_normals=raw_data.vertex_normals,
            faces=raw_data.faces,
            face_normals=raw_data.face_normals,
            joints=raw_data.joints,
            tails=raw_data.tails,
            skin=raw_data.skin,
            no_skin=raw_data.no_skin,
            parents=raw_data.parents,
            names=raw_data.names,
            matrix_local=raw_data.matrix_local,
            meta={},
        )
    
    def get_tokenize_input(self) -> TokenizeInput:
        children = defaultdict(list)
        
        for (id, p) in enumerate(self.parents):
            if p is not None:
                children[p].append(id)
        bones = []
        branch = []
        is_leaf = []
        last = None
        for i in range(self.J):
            is_leaf.append(len(children[i])==0)
            if i == 0:
                bones.append(np.concatenate([self.joints[i], self.joints[i]]))
                branch.append(False)
            else:
                pid = self.parents[i]
                bones.append(np.concatenate([self.joints[pid], self.joints[i]]))
                branch.append(pid!=last)
            last = i
        bones = np.stack(bones)
        branch = np.array(branch, dtype=bool)
        is_leaf = np.array(is_leaf, dtype=bool)
        return TokenizeInput(
            bones=bones,
            tails=self.tails,
            branch=branch,
            is_leaf=is_leaf,
            no_skin=self.no_skin,
            cls=self.cls,
            parts_bias=self.parts_bias,
        )
    
    def export_pc(self, path: str, with_normal: bool=True, normal_size=0.01):
        '''
        export point cloud
        '''
        vertices    = self.vertices
        normals     = self.vertex_normals
        if self.sampled_vertices is not None:
            vertices    = self.sampled_vertices
            normals     = self.sampled_normals
        if with_normal == False:
            normals = None
        self._export_pc(vertices=vertices, path=path, vertex_normals=normals, normal_size=normal_size)
    
    def export_mesh(self, path: str):
        '''
        export mesh
        '''
        self._export_mesh(vertices=self.vertices, faces=self.faces, path=path)
    
    def export_skeleton(self, path: str):
        '''
        export spring
        '''
        self._export_skeleton(joints=self.joints, parents=self.parents, path=path)
    
    def export_skeleton_sequence(self, path: str):
        '''
        export spring
        '''
        self._export_skeleton_sequence(joints=self.joints, parents=self.parents, path=path)
    
    def export_fbx(
        self,
        path: str,
        vertex_group_name: str,
        extrude_size: float=0.03,
        group_per_vertex: int=-1,
        add_root: bool=False,
        do_not_normalize: bool=False,
        use_extrude_bone: bool=True,
        use_connect_unique_child: bool=True,
        extrude_from_parent: bool=True,
        use_tail: bool=False,
        use_origin: bool=False,
    ):
        '''
        export the whole model with skining
        '''
        self._export_fbx(
            path=path,
            vertices=self.vertices if use_origin else self.sampled_vertices,
            joints=self.joints,
            skin=self.sampled_vertex_groups[vertex_group_name],
            parents=self.parents,
            names=self.names,
            faces=self.faces if use_origin else None,
            extrude_size=extrude_size,
            group_per_vertex=group_per_vertex,
            add_root=add_root,
            do_not_normalize=do_not_normalize,
            use_extrude_bone=use_extrude_bone,
            use_connect_unique_child=use_connect_unique_child,
            extrude_from_parent=extrude_from_parent,
            tails=self.tails if use_tail else None,
        )

    def export_render(self, path: str, resolution: Tuple[int, int]=[256, 256], use_tail: bool=False):
        if use_tail:
            assert self.tails is not None
            self._export_render(
                path=path,
                vertices=self.vertices,
                faces=self.faces,
                bones=np.concatenate([self.joints, self.tails], axis=-1),
                resolution=resolution,
            )
        else:
            pjoints = self.joints[self.parents[1:]]
            self._export_render(
                path=path,
                vertices=self.vertices,
                faces=self.faces,
                bones=np.concatenate([pjoints, self.joints[1:]], axis=-1),
                resolution=resolution,
            )