File size: 14,135 Bytes
f499d3b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
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,
            )