File size: 3,484 Bytes
2252f3d |
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 |
# -*- coding: utf-8 -*-
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed
# a license agreement with MPG or you get the right to use the computer
# program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and
# liable to prosecution.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de
from typing import NewType, Union, Optional
from dataclasses import dataclass, asdict, fields
import numpy as np
import torch
Tensor = NewType('Tensor', torch.Tensor)
Array = NewType('Array', np.ndarray)
@dataclass
class ModelOutput:
vertices: Optional[Tensor] = None
joints: Optional[Tensor] = None
full_pose: Optional[Tensor] = None
global_orient: Optional[Tensor] = None
transl: Optional[Tensor] = None
def __getitem__(self, key):
return getattr(self, key)
def get(self, key, default=None):
return getattr(self, key, default)
def __iter__(self):
return self.keys()
def keys(self):
keys = [t.name for t in fields(self)]
return iter(keys)
def values(self):
values = [getattr(self, t.name) for t in fields(self)]
return iter(values)
def items(self):
data = [(t.name, getattr(self, t.name)) for t in fields(self)]
return iter(data)
@dataclass
class SMPLOutput(ModelOutput):
betas: Optional[Tensor] = None
body_pose: Optional[Tensor] = None
@dataclass
class SMPLHOutput(SMPLOutput):
left_hand_pose: Optional[Tensor] = None
right_hand_pose: Optional[Tensor] = None
transl: Optional[Tensor] = None
@dataclass
class SMPLXOutput(SMPLHOutput):
expression: Optional[Tensor] = None
jaw_pose: Optional[Tensor] = None
joint_transformation: Optional[Tensor] = None
vertex_transformation: Optional[Tensor] = None
@dataclass
class MANOOutput(ModelOutput):
betas: Optional[Tensor] = None
hand_pose: Optional[Tensor] = None
@dataclass
class FLAMEOutput(ModelOutput):
betas: Optional[Tensor] = None
expression: Optional[Tensor] = None
jaw_pose: Optional[Tensor] = None
neck_pose: Optional[Tensor] = None
def find_joint_kin_chain(joint_id, kinematic_tree):
kin_chain = []
curr_idx = joint_id
while curr_idx != -1:
kin_chain.append(curr_idx)
curr_idx = kinematic_tree[curr_idx]
return kin_chain
def to_tensor(array: Union[Array, Tensor], dtype=torch.float32) -> Tensor:
if torch.is_tensor(array):
return array
else:
return torch.tensor(array, dtype=dtype)
class Struct(object):
def __init__(self, **kwargs):
for key, val in kwargs.items():
setattr(self, key, val)
def to_np(array, dtype=np.float32):
if 'scipy.sparse' in str(type(array)):
array = array.todense()
return np.array(array, dtype=dtype)
def rot_mat_to_euler(rot_mats):
# Calculates rotation matrix to euler angles
# Careful for extreme cases of eular angles like [0.0, pi, 0.0]
sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
return torch.atan2(-rot_mats[:, 2, 0], sy)
|