Spaces:
Running
on
Zero
Running
on
Zero
from dataclasses import dataclass | |
import numpy as np | |
from numpy import ndarray | |
import os | |
from typing import Union, List, Tuple | |
from .exporter import Exporter | |
from ..tokenizer.spec import DetokenzeOutput | |
from .order import Order | |
class RawData(Exporter): | |
''' | |
Dataclass to handle data from processed model files. | |
''' | |
# vertices of the mesh, shape (N, 3), float32 | |
vertices: Union[ndarray, None] | |
# normals of vertices, shape (N, 3), float32 | |
vertex_normals: Union[ndarray, None] | |
# faces of mesh, shape (F, 3), face id starts from 0 to F-1, int64 | |
faces: Union[ndarray, None] | |
# face normal of mesh, shape (F, 3), float32 | |
face_normals: Union[ndarray, None] | |
# joints of bones, shape (J, 3), float32 | |
joints: Union[ndarray, None] | |
# tails of joints, shape (J, 3), float32 | |
tails: Union[ndarray, None] | |
# skinning of joints, shape (N, J), float32 | |
skin: Union[ndarray, None] | |
# whether the joint has skin, bool | |
no_skin: Union[ndarray, None] | |
# parents of joints, None represents no parent(a root joint) | |
# make sure parent[k] < k | |
parents: Union[List[Union[int, None]], None] | |
# names of joints | |
names: Union[List[str], None] | |
# local coordinate | |
matrix_local: Union[ndarray, None] | |
# path to data | |
path: Union[str, None]=None | |
# data cls | |
cls: Union[str, None]=None | |
def load(path: str) -> 'RawData': | |
data = np.load(path, allow_pickle=True) | |
d = {name: data[name][()] for name in data} | |
d['path'] = path | |
return RawData(**d) | |
def save(self, path: str): | |
os.makedirs(os.path.dirname(path), exist_ok=True) | |
np.savez(file=path, **self.__dict__) | |
def N(self): | |
''' | |
number of vertices | |
''' | |
return self.vertices.shape[0] | |
def F(self): | |
''' | |
number of faces | |
''' | |
return self.faces.shape[0] | |
def J(self): | |
''' | |
number of joints | |
''' | |
return self.joints.shape[0] | |
def check(self): | |
if self.names is not None and self.joints is not None: | |
assert len(self.names) == self.J | |
if self.names is not None and self.parents is not None: | |
assert len(self.names) == len(self.parents) | |
if self.parents is not None: | |
for (i, pid) in enumerate(self.parents): | |
if i==0: | |
assert pid is None | |
else: | |
assert pid is not None | |
assert pid < i | |
def export_pc(self, path: str, with_normal: bool=True, normal_size=0.01): | |
''' | |
export point cloud | |
''' | |
if with_normal: | |
self._export_pc(vertices=self.vertices, path=path, vertex_normals=self.vertex_normals, normal_size=normal_size) | |
else: | |
self._export_pc(vertices=self.vertices, path=path, vertex_normals=None, 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, | |
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, | |
custom_vertex_group: Union[ndarray, None]=None, | |
): | |
''' | |
export the whole model with skining | |
''' | |
self._export_fbx( | |
path=path, | |
vertices=self.vertices, | |
joints=self.joints, | |
skin=self.skin if custom_vertex_group is None else custom_vertex_group, | |
parents=self.parents, | |
names=self.names, | |
faces=self.faces, | |
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]): | |
self._export_render( | |
path=path, | |
vertices=self.vertices, | |
faces=self.faces, | |
bones=np.concatenate([self.joints, self.tails], axis=-1), | |
resolution=resolution, | |
) | |
class RawSkeleton(Exporter): | |
''' | |
Dataclass to handle skeleton from AR. | |
''' | |
# joints of bones, shape (J, 3), float32 | |
joints: Union[ndarray, None] | |
# tails of joints, shape (J, 3), float32 | |
tails: Union[ndarray, None] | |
# whether the joint has skin, bool | |
no_skin: Union[ndarray, None] | |
# parents of joints, None represents no parent(a root joint) | |
# make sure parent[k] < k | |
parents: Union[List[Union[int, None]], None] | |
# names of joints | |
names: Union[List[str], None] | |
def load(path: str) -> 'RawSkeleton': | |
data = np.load(path, allow_pickle=True) | |
return RawSkeleton(**{name: data[name][()] for name in data}) | |
def save(self, path: str): | |
os.makedirs(os.path.dirname(path), exist_ok=True) | |
np.savez(file=path, **self.__dict__) | |
def from_detokenize_output(res: DetokenzeOutput, order: Union[Order, None]) -> 'RawSkeleton': | |
J = len(res.bones) | |
names = order.make_names(cls=res.cls, parts=res.parts, num_bones=J) | |
joints = res.joints | |
p_joints = res.p_joints | |
parents = [] | |
for (i, joint) in enumerate(joints): | |
if i == 0: | |
parents.append(None) | |
continue | |
p_joint = p_joints[i] | |
dis = 999999 | |
pid = None | |
for j in reversed(range(i)): | |
n_dis = ((joints[j] - p_joint)**2).sum() | |
if n_dis < dis: | |
pid = j | |
dis = n_dis | |
parents.append(pid) | |
return RawSkeleton( | |
joints=joints, | |
tails=res.tails, | |
no_skin=res.no_skin, | |
parents=parents, | |
names=names, | |
) | |
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, | |
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, | |
): | |
''' | |
export the whole model with skining | |
''' | |
self._export_fbx( | |
path=path, | |
vertices=None, | |
joints=self.joints, | |
skin=None, | |
parents=self.parents, | |
names=self.names, | |
faces=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]): | |
self._export_render( | |
path=path, | |
vertices=None, | |
faces=None, | |
bones=np.concatenate([self.joints, self.tails], axis=-1), | |
resolution=resolution, | |
) | |
class RawSkin(Exporter): | |
''' | |
Dataclass to handle skeleton from AR. | |
''' | |
# skin, shape (J, N) | |
skin: ndarray | |
# always sampled, shape (N, 3) | |
vertices: Union[ndarray, None]=None | |
# for future use, shape (J, 3) | |
joints: Union[ndarray, None]=None | |
def load(path: str) -> 'RawSkin': | |
data = np.load(path, allow_pickle=True) | |
return RawSkin(**{name: data[name][()] for name in data}) | |
def save(self, path: str): | |
os.makedirs(os.path.dirname(path), exist_ok=True) | |
np.savez(file=path, **self.__dict__) |