Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,450 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 |
from dataclasses import dataclass
from typing import Tuple, Union, List, Dict
from numpy import ndarray
import numpy as np
from abc import ABC, abstractmethod
from scipy.spatial.transform import Rotation as R
from .spec import ConfigSpec
from .asset import Asset
from .utils import axis_angle_to_matrix
@dataclass(frozen=True)
class AugmentAffineConfig(ConfigSpec):
# final normalization cube
normalize_into: Tuple[float, float]
# randomly scale coordinates with probability p
random_scale_p: float
# scale range (lower, upper)
random_scale: Tuple[float, float]
# randomly shift coordinates with probability p
random_shift_p: float
# shift range (lower, upper)
random_shift: Tuple[float, float]
@classmethod
def parse(cls, config) -> Union['AugmentAffineConfig', None]:
if config is None:
return None
cls.check_keys(config)
return AugmentAffineConfig(
normalize_into=config.normalize_into,
random_scale_p=config.get('random_scale_p', 0.),
random_scale=config.get('random_scale', [1., 1.]),
random_shift_p=config.get('random_shift_p', 0.),
random_shift=config.get('random_shift', [0., 0.]),
)
@dataclass(frozen=True)
class AugmentConfig(ConfigSpec):
'''
Config to handle final easy augmentation of vertices, normals and bones before sampling.
'''
augment_affine_config: Union[AugmentAffineConfig, None]
@classmethod
def parse(cls, config) -> 'AugmentConfig':
cls.check_keys(config)
return AugmentConfig(
augment_affine_config=AugmentAffineConfig.parse(config.get('augment_affine_config', None)),
)
class Augment(ABC):
'''
Abstract class for augmentation
'''
def __init__(self):
pass
@abstractmethod
def transform(self, asset: Asset, **kwargs):
pass
@abstractmethod
def inverse(self, asset: Asset):
pass
class AugmentAffine(Augment):
def __init__(self, config: AugmentAffineConfig):
super().__init__()
self.config = config
def _apply(self, v: ndarray, trans: ndarray) -> ndarray:
return np.matmul(v, trans[:3, :3].transpose()) + trans[:3, 3]
def transform(self, asset: Asset, **kwargs):
bound_min = asset.vertices.min(axis=0)
bound_max = asset.vertices.max(axis=0)
if asset.joints is not None:
joints_bound_min = asset.joints.min(axis=0)
joints_bound_max = asset.joints.max(axis=0)
bound_min = np.minimum(bound_min, joints_bound_min)
bound_max = np.maximum(bound_max, joints_bound_max)
trans_vertex = np.eye(4, dtype=np.float32)
trans_vertex = _trans_to_m(-(bound_max + bound_min)/2) @ trans_vertex
# scale into the cube
normalize_into = self.config.normalize_into
scale = np.max((bound_max - bound_min) / (normalize_into[1] - normalize_into[0]))
trans_vertex = _scale_to_m(1. / scale) @ trans_vertex
bias = (normalize_into[0] + normalize_into[1]) / 2
trans_vertex = _trans_to_m(np.array([bias, bias, bias], dtype=np.float32)) @ trans_vertex
if np.random.rand() < self.config.random_scale_p:
scale = _scale_to_m(np.random.uniform(self.config.random_scale[0], self.config.random_scale[1]))
trans_vertex = scale @ trans_vertex
if np.random.rand() < self.config.random_shift_p:
l, r = self.config.random_shift
shift = _trans_to_m(np.array([np.random.uniform(l, r), np.random.uniform(l, r), np.random.uniform(l, r)]), dtype=np.float32)
trans_vertex = shift @ trans_vertex
asset.vertices = self._apply(asset.vertices, trans_vertex)
# do not affect scale in matrix
if asset.matrix_local is not None:
asset.matrix_local[:, :, 3:4] = trans_vertex @ asset.matrix_local[:, :, 3:4]
if asset.pose_matrix is not None:
asset.pose_matrix[:, :, 3:4] = trans_vertex @ asset.pose_matrix[:, :, 3:4]
# do not affect normal here
if asset.joints is not None:
asset.joints = self._apply(asset.joints, trans_vertex)
if asset.tails is not None:
asset.tails = self._apply(asset.tails, trans_vertex)
self.trans_vertex = trans_vertex
def inverse(self, asset: Asset):
m = np.linalg.inv(self.trans_vertex)
asset.vertices = self._apply(asset.vertices, m)
if asset.joints is not None:
asset.joints = self._apply(asset.joints, m)
if asset.tails is not None:
asset.tails = self._apply(asset.tails, m)
def _trans_to_m(v: ndarray):
m = np.eye(4, dtype=np.float32)
m[0:3, 3] = v
return m
def _scale_to_m(r: ndarray):
m = np.zeros((4, 4), dtype=np.float32)
m[0, 0] = r
m[1, 1] = r
m[2, 2] = r
m[3, 3] = 1.
return m
def get_augments(config: AugmentConfig) -> Tuple[List[Augment], List[Augment]]:
first_augments = [] # augments before sample
second_augments = [] # augments after sample
augment_affine_config = config.augment_affine_config
if augment_affine_config is not None:
second_augments.append(AugmentAffine(config=augment_affine_config))
return first_augments, second_augments |