|
import numpy as np |
|
import scipy.linalg as linalg |
|
|
|
from visualization import Animation |
|
from visualization import AnimationStructure |
|
|
|
from visualization.Quaternions import Quaternions |
|
|
|
|
|
class BasicInverseKinematics: |
|
""" |
|
Basic Inverse Kinematics Solver |
|
|
|
This is an extremely simple full body IK |
|
solver. |
|
|
|
It works given the following conditions: |
|
|
|
* All joint targets must be specified |
|
* All joint targets must be in reach |
|
* All joint targets must not differ |
|
extremely from the starting pose |
|
* No bone length constraints can be violated |
|
* The root translation and rotation are |
|
set to good initial values |
|
|
|
It works under the observation that if the |
|
_directions_ the joints are pointing toward |
|
match the _directions_ of the vectors between |
|
the target joints then the pose should match |
|
that of the target pose. |
|
|
|
Therefore it iterates over joints rotating |
|
each joint such that the vectors between it |
|
and it's children match that of the target |
|
positions. |
|
|
|
Parameters |
|
---------- |
|
|
|
animation : Animation |
|
animation input |
|
|
|
positions : (F, J, 3) ndarray |
|
target positions for each frame F |
|
and each joint J |
|
|
|
iterations : int |
|
Optional number of iterations. |
|
If the above conditions are met |
|
1 iteration should be enough, |
|
therefore the default is 1 |
|
|
|
silent : bool |
|
Optional if to suppress output |
|
defaults to False |
|
""" |
|
|
|
def __init__(self, animation, positions, iterations=1, silent=True): |
|
|
|
self.animation = animation |
|
self.positions = positions |
|
self.iterations = iterations |
|
self.silent = silent |
|
|
|
def __call__(self): |
|
|
|
children = AnimationStructure.children_list(self.animation.parents) |
|
|
|
for i in range(self.iterations): |
|
|
|
for j in AnimationStructure.joints(self.animation.parents): |
|
|
|
c = np.array(children[j]) |
|
if len(c) == 0: continue |
|
|
|
anim_transforms = Animation.transforms_global(self.animation) |
|
anim_positions = anim_transforms[:, :, :3, 3] |
|
anim_rotations = Quaternions.from_transforms(anim_transforms) |
|
|
|
jdirs = anim_positions[:, c] - anim_positions[:, np.newaxis, j] |
|
ddirs = self.positions[:, c] - anim_positions[:, np.newaxis, j] |
|
|
|
jsums = np.sqrt(np.sum(jdirs ** 2.0, axis=-1)) + 1e-10 |
|
dsums = np.sqrt(np.sum(ddirs ** 2.0, axis=-1)) + 1e-10 |
|
|
|
jdirs = jdirs / jsums[:, :, np.newaxis] |
|
ddirs = ddirs / dsums[:, :, np.newaxis] |
|
|
|
angles = np.arccos(np.sum(jdirs * ddirs, axis=2).clip(-1, 1)) |
|
axises = np.cross(jdirs, ddirs) |
|
axises = -anim_rotations[:, j, np.newaxis] * axises |
|
|
|
rotations = Quaternions.from_angle_axis(angles, axises) |
|
|
|
if rotations.shape[1] == 1: |
|
averages = rotations[:, 0] |
|
else: |
|
averages = Quaternions.exp(rotations.log().mean(axis=-2)) |
|
|
|
self.animation.rotations[:, j] = self.animation.rotations[:, j] * averages |
|
|
|
if not self.silent: |
|
anim_positions = Animation.positions_global(self.animation) |
|
error = np.mean(np.sum((anim_positions - self.positions) ** 2.0, axis=-1) ** 0.5) |
|
print('[BasicInverseKinematics] Iteration %i Error: %f' % (i + 1, error)) |
|
|
|
return self.animation |
|
|
|
|
|
class JacobianInverseKinematics: |
|
""" |
|
Jacobian Based Full Body IK Solver |
|
|
|
This is a full body IK solver which |
|
uses the dampened least squares inverse |
|
jacobian method. |
|
|
|
It should remain fairly stable and effective |
|
even for joint positions which are out of |
|
reach and it can also take any number of targets |
|
to treat as end effectors. |
|
|
|
Parameters |
|
---------- |
|
|
|
animation : Animation |
|
animation to solve inverse problem on |
|
|
|
targets : {int : (F, 3) ndarray} |
|
Dictionary of target positions for each |
|
frame F, mapping joint index to |
|
a target position |
|
|
|
references : (F, 3) |
|
Optional list of J joint position |
|
references for which the result |
|
should bias toward |
|
|
|
iterations : int |
|
Optional number of iterations to |
|
compute. More iterations results in |
|
better accuracy but takes longer to |
|
compute. Default is 10. |
|
|
|
recalculate : bool |
|
Optional if to recalcuate jacobian |
|
each iteration. Gives better accuracy |
|
but slower to compute. Defaults to True |
|
|
|
damping : float |
|
Optional damping constant. Higher |
|
damping increases stability but |
|
requires more iterations to converge. |
|
Defaults to 5.0 |
|
|
|
secondary : float |
|
Force, or bias toward secondary target. |
|
Defaults to 0.25 |
|
|
|
silent : bool |
|
Optional if to suppress output |
|
defaults to False |
|
""" |
|
|
|
def __init__(self, animation, targets, |
|
references=None, iterations=10, |
|
recalculate=True, damping=2.0, |
|
secondary=0.25, translate=False, |
|
silent=False, weights=None, |
|
weights_translate=None): |
|
|
|
self.animation = animation |
|
self.targets = targets |
|
self.references = references |
|
|
|
self.iterations = iterations |
|
self.recalculate = recalculate |
|
self.damping = damping |
|
self.secondary = secondary |
|
self.translate = translate |
|
self.silent = silent |
|
self.weights = weights |
|
self.weights_translate = weights_translate |
|
|
|
def cross(self, a, b): |
|
o = np.empty(b.shape) |
|
o[..., 0] = a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1] |
|
o[..., 1] = a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2] |
|
o[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] |
|
return o |
|
|
|
def jacobian(self, x, fp, fr, ts, dsc, tdsc): |
|
|
|
""" Find parent rotations """ |
|
prs = fr[:, self.animation.parents] |
|
prs[:, 0] = Quaternions.id((1)) |
|
|
|
""" Find global positions of target joints """ |
|
tps = fp[:, np.array(list(ts.keys()))] |
|
|
|
""" Get partial rotations """ |
|
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1] * 3:3], np.array([[[0, 1, 0]]])) |
|
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1] * 3:3], np.array([[[0, 0, 1]]])) |
|
|
|
""" Find axis of rotations """ |
|
es = np.empty((len(x), fr.shape[1] * 3, 3)) |
|
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]]) |
|
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]])) |
|
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]]))) |
|
|
|
""" Construct Jacobian """ |
|
j = fp.repeat(3, axis=1) |
|
j = dsc[np.newaxis, :, :, np.newaxis] * (tps[:, np.newaxis, :] - j[:, :, np.newaxis]) |
|
j = self.cross(es[:, :, np.newaxis, :], j) |
|
j = np.swapaxes(j.reshape((len(x), fr.shape[1] * 3, len(ts) * 3)), 1, 2) |
|
|
|
if self.translate: |
|
es = np.empty((len(x), fr.shape[1] * 3, 3)) |
|
es[:, 0::3] = prs * np.array([[[1, 0, 0]]]) |
|
es[:, 1::3] = prs * np.array([[[0, 1, 0]]]) |
|
es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) |
|
|
|
jt = tdsc[np.newaxis, :, :, np.newaxis] * es[:, :, np.newaxis, :].repeat(tps.shape[1], axis=2) |
|
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1] * 3, len(ts) * 3)), 1, 2) |
|
|
|
j = np.concatenate([j, jt], axis=-1) |
|
|
|
return j |
|
|
|
|
|
def __call__(self, descendants=None, gamma=1.0): |
|
|
|
self.descendants = descendants |
|
|
|
""" Calculate Masses """ |
|
if self.weights is None: |
|
self.weights = np.ones(self.animation.shape[1]) |
|
|
|
if self.weights_translate is None: |
|
self.weights_translate = np.ones(self.animation.shape[1]) |
|
|
|
""" Calculate Descendants """ |
|
if self.descendants is None: |
|
self.descendants = AnimationStructure.descendants_mask(self.animation.parents) |
|
|
|
self.tdescendants = np.eye(self.animation.shape[1]) + self.descendants |
|
|
|
self.first_descendants = self.descendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype(int) |
|
self.first_tdescendants = self.tdescendants[:, np.array(list(self.targets.keys()))].repeat(3, axis=0).astype( |
|
int) |
|
|
|
""" Calculate End Effectors """ |
|
self.endeff = np.array(list(self.targets.values())) |
|
self.endeff = np.swapaxes(self.endeff, 0, 1) |
|
|
|
if not self.references is None: |
|
self.second_descendants = self.descendants.repeat(3, axis=0).astype(int) |
|
self.second_tdescendants = self.tdescendants.repeat(3, axis=0).astype(int) |
|
self.second_targets = dict([(i, self.references[:, i]) for i in range(self.references.shape[1])]) |
|
|
|
nf = len(self.animation) |
|
nj = self.animation.shape[1] |
|
|
|
if not self.silent: |
|
gp = Animation.positions_global(self.animation) |
|
gp = gp[:, np.array(list(self.targets.keys()))] |
|
error = np.mean(np.sqrt(np.sum((self.endeff - gp) ** 2.0, axis=2))) |
|
print('[JacobianInverseKinematics] Start | Error: %f' % error) |
|
|
|
for i in range(self.iterations): |
|
|
|
""" Get Global Rotations & Positions """ |
|
gt = Animation.transforms_global(self.animation) |
|
gp = gt[:, :, :, 3] |
|
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis] |
|
gr = Quaternions.from_transforms(gt) |
|
|
|
x = self.animation.rotations.euler().reshape(nf, -1) |
|
w = self.weights.repeat(3) |
|
|
|
if self.translate: |
|
x = np.hstack([x, self.animation.positions.reshape(nf, -1)]) |
|
w = np.hstack([w, self.weights_translate.repeat(3)]) |
|
|
|
""" Generate Jacobian """ |
|
if self.recalculate or i == 0: |
|
j = self.jacobian(x, gp, gr, self.targets, self.first_descendants, self.first_tdescendants) |
|
|
|
""" Update Variables """ |
|
l = self.damping * (1.0 / (w + 0.001)) |
|
d = (l * l) * np.eye(x.shape[1]) |
|
e = gamma * (self.endeff.reshape(nf, -1) - gp[:, np.array(list(self.targets.keys()))].reshape(nf, -1)) |
|
|
|
x += np.array(list(map(lambda jf, ef: |
|
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e))) |
|
|
|
""" Generate Secondary Jacobian """ |
|
if self.references is not None: |
|
|
|
ns = np.array(list(map(lambda jf: |
|
np.eye(x.shape[1]) - linalg.solve(jf.T.dot(jf) + d, jf.T.dot(jf)), j))) |
|
|
|
if self.recalculate or i == 0: |
|
j2 = self.jacobian(x, gp, gr, self.second_targets, self.second_descendants, |
|
self.second_tdescendants) |
|
|
|
e2 = self.secondary * (self.references.reshape(nf, -1) - gp.reshape(nf, -1)) |
|
|
|
x += np.array(list(map(lambda nsf, j2f, e2f: |
|
nsf.dot(linalg.lu_solve(linalg.lu_factor(j2f.T.dot(j2f) + d), j2f.T.dot(e2f))), |
|
ns, j2, e2))) |
|
|
|
""" Set Back Rotations / Translations """ |
|
self.animation.rotations = Quaternions.from_euler( |
|
x[:, :nj * 3].reshape((nf, nj, 3)), order='xyz', world=True) |
|
|
|
if self.translate: |
|
self.animation.positions = x[:, nj * 3:].reshape((nf, nj, 3)) |
|
|
|
""" Generate Error """ |
|
|
|
if not self.silent: |
|
gp = Animation.positions_global(self.animation) |
|
gp = gp[:, np.array(list(self.targets.keys()))] |
|
error = np.mean(np.sum((self.endeff - gp) ** 2.0, axis=2) ** 0.5) |
|
print('[JacobianInverseKinematics] Iteration %i | Error: %f' % (i + 1, error)) |
|
return self.animation |
|
|
|
|
|
class BasicJacobianIK: |
|
""" |
|
Same interface as BasicInverseKinematics |
|
but uses the Jacobian IK Solver Instead |
|
""" |
|
|
|
def __init__(self, animation, positions, iterations=10, silent=True, **kw): |
|
targets = dict([(i, positions[:, i]) for i in range(positions.shape[1])]) |
|
self.ik = JacobianInverseKinematics(animation, targets, iterations=iterations, silent=silent, **kw) |
|
|
|
def __call__(self, **kw): |
|
return self.ik(**kw) |
|
|
|
|
|
class ICP: |
|
|
|
def __init__(self, |
|
anim, rest, weights, mesh, goal, |
|
find_closest=True, damping=10, |
|
iterations=10, silent=True, |
|
translate=True, recalculate=True, |
|
weights_translate=None): |
|
|
|
self.animation = anim |
|
self.rest = rest |
|
self.vweights = weights |
|
self.mesh = mesh |
|
self.goal = goal |
|
self.find_closest = find_closest |
|
self.iterations = iterations |
|
self.silent = silent |
|
self.translate = translate |
|
self.damping = damping |
|
self.weights = None |
|
self.weights_translate = weights_translate |
|
self.recalculate = recalculate |
|
|
|
def cross(self, a, b): |
|
o = np.empty(b.shape) |
|
o[..., 0] = a[..., 1] * b[..., 2] - a[..., 2] * b[..., 1] |
|
o[..., 1] = a[..., 2] * b[..., 0] - a[..., 0] * b[..., 2] |
|
o[..., 2] = a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] |
|
return o |
|
|
|
def jacobian(self, x, fp, fr, goal, weights, des_r, des_t): |
|
|
|
""" Find parent rotations """ |
|
prs = fr[:, self.animation.parents] |
|
prs[:, 0] = Quaternions.id((1)) |
|
|
|
""" Get partial rotations """ |
|
qys = Quaternions.from_angle_axis(x[:, 1:prs.shape[1] * 3:3], np.array([[[0, 1, 0]]])) |
|
qzs = Quaternions.from_angle_axis(x[:, 2:prs.shape[1] * 3:3], np.array([[[0, 0, 1]]])) |
|
|
|
""" Find axis of rotations """ |
|
es = np.empty((len(x), fr.shape[1] * 3, 3)) |
|
es[:, 0::3] = ((prs * qzs) * qys) * np.array([[[1, 0, 0]]]) |
|
es[:, 1::3] = ((prs * qzs) * np.array([[[0, 1, 0]]])) |
|
es[:, 2::3] = ((prs * np.array([[[0, 0, 1]]]))) |
|
|
|
""" Construct Jacobian """ |
|
j = fp.repeat(3, axis=1) |
|
j = des_r[np.newaxis, :, :, :, np.newaxis] * ( |
|
goal[:, np.newaxis, :, np.newaxis] - j[:, :, np.newaxis, np.newaxis]) |
|
j = np.sum(j * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3) |
|
j = self.cross(es[:, :, np.newaxis, :], j) |
|
j = np.swapaxes(j.reshape((len(x), fr.shape[1] * 3, goal.shape[1] * 3)), 1, 2) |
|
|
|
if self.translate: |
|
es = np.empty((len(x), fr.shape[1] * 3, 3)) |
|
es[:, 0::3] = prs * np.array([[[1, 0, 0]]]) |
|
es[:, 1::3] = prs * np.array([[[0, 1, 0]]]) |
|
es[:, 2::3] = prs * np.array([[[0, 0, 1]]]) |
|
|
|
jt = des_t[np.newaxis, :, :, :, np.newaxis] * es[:, :, np.newaxis, np.newaxis, :].repeat(goal.shape[1], |
|
axis=2) |
|
jt = np.sum(jt * weights[np.newaxis, np.newaxis, :, :, np.newaxis], 3) |
|
jt = np.swapaxes(jt.reshape((len(x), fr.shape[1] * 3, goal.shape[1] * 3)), 1, 2) |
|
|
|
j = np.concatenate([j, jt], axis=-1) |
|
|
|
return j |
|
|
|
|
|
def __call__(self, descendants=None, maxjoints=4, gamma=1.0, transpose=False): |
|
|
|
""" Calculate Masses """ |
|
if self.weights is None: |
|
self.weights = np.ones(self.animation.shape[1]) |
|
|
|
if self.weights_translate is None: |
|
self.weights_translate = np.ones(self.animation.shape[1]) |
|
|
|
nf = len(self.animation) |
|
nj = self.animation.shape[1] |
|
nv = self.goal.shape[1] |
|
|
|
weightids = np.argsort(-self.vweights, axis=1)[:, :maxjoints] |
|
weightvls = np.array(list(map(lambda w, i: w[i], self.vweights, weightids))) |
|
weightvls = weightvls / weightvls.sum(axis=1)[..., np.newaxis] |
|
|
|
if descendants is None: |
|
self.descendants = AnimationStructure.descendants_mask(self.animation.parents) |
|
else: |
|
self.descendants = descendants |
|
|
|
des_r = np.eye(nj) + self.descendants |
|
des_r = des_r[:, weightids].repeat(3, axis=0) |
|
|
|
des_t = np.eye(nj) + self.descendants |
|
des_t = des_t[:, weightids].repeat(3, axis=0) |
|
|
|
if not self.silent: |
|
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints) |
|
error = np.mean(np.sqrt(np.sum((curr - self.goal) ** 2.0, axis=-1))) |
|
print('[ICP] Start | Error: %f' % error) |
|
|
|
for i in range(self.iterations): |
|
|
|
""" Get Global Rotations & Positions """ |
|
gt = Animation.transforms_global(self.animation) |
|
gp = gt[:, :, :, 3] |
|
gp = gp[:, :, :3] / gp[:, :, 3, np.newaxis] |
|
gr = Quaternions.from_transforms(gt) |
|
|
|
x = self.animation.rotations.euler().reshape(nf, -1) |
|
w = self.weights.repeat(3) |
|
|
|
if self.translate: |
|
x = np.hstack([x, self.animation.positions.reshape(nf, -1)]) |
|
w = np.hstack([w, self.weights_translate.repeat(3)]) |
|
|
|
""" Get Current State """ |
|
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh, maxjoints=maxjoints) |
|
|
|
""" Find Cloest Points """ |
|
if self.find_closest: |
|
mapping = np.argmin( |
|
(curr[:, :, np.newaxis] - |
|
self.goal[:, np.newaxis, :]) ** 2.0, axis=2) |
|
e = gamma * (np.array(list(map(lambda g, m: g[m], self.goal, mapping))) - curr).reshape(nf, -1) |
|
else: |
|
e = gamma * (self.goal - curr).reshape(nf, -1) |
|
|
|
""" Generate Jacobian """ |
|
if self.recalculate or i == 0: |
|
j = self.jacobian(x, gp, gr, self.goal, weightvls, des_r, des_t) |
|
|
|
""" Update Variables """ |
|
l = self.damping * (1.0 / (w + 1e-10)) |
|
d = (l * l) * np.eye(x.shape[1]) |
|
|
|
if transpose: |
|
x += np.array(list(map(lambda jf, ef: jf.T.dot(ef), j, e))) |
|
else: |
|
x += np.array(list(map(lambda jf, ef: |
|
linalg.lu_solve(linalg.lu_factor(jf.T.dot(jf) + d), jf.T.dot(ef)), j, e))) |
|
|
|
""" Set Back Rotations / Translations """ |
|
self.animation.rotations = Quaternions.from_euler( |
|
x[:, :nj * 3].reshape((nf, nj, 3)), order='xyz', world=True) |
|
|
|
if self.translate: |
|
self.animation.positions = x[:, nj * 3:].reshape((nf, nj, 3)) |
|
|
|
if not self.silent: |
|
curr = Animation.skin(self.animation, self.rest, self.vweights, self.mesh) |
|
error = np.mean(np.sqrt(np.sum((curr - self.goal) ** 2.0, axis=-1))) |
|
print('[ICP] Iteration %i | Error: %f' % (i + 1, error)) |
|
|
|
import torch |
|
from torch import nn |
|
class InverseKinematics: |
|
def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains): |
|
self.rotations = rotations.cuda() |
|
self.rotations.requires_grad_(True) |
|
self.position = positions.cuda() |
|
self.position.requires_grad_(True) |
|
|
|
self.parents = parents |
|
self.offset = offset.cuda() |
|
self.constrains = constrains.cuda() |
|
|
|
self.optimizer = torch.optim.AdamW([self.position, self.rotations], lr=5e-2, betas=(0.9, 0.999)) |
|
self.crit = nn.MSELoss() |
|
self.weights = torch.ones([1,22,1]).cuda() |
|
self.weights[:, [4, 8]] = 0.8 |
|
self.weights[:, [1, 5]] = 2. |
|
|
|
def step(self): |
|
self.optimizer.zero_grad() |
|
glb = self.forward(self.rotations, self.position, self.offset, order='', quater=True, world=True) |
|
|
|
loss = self.crit(glb*self.weights, self.constrains*self.weights) |
|
|
|
loss += 0.5 * self.crit(self.rotations[1:, [3, 7, 12, 16, 20]], self.rotations[:-1, [3, 7, 12, 16, 20]]) + 0.1 * self.crit(self.rotations[1:], self.rotations[:-1]) |
|
loss.backward() |
|
self.optimizer.step() |
|
self.glb = glb |
|
return loss.item() |
|
|
|
def tloss(self, time): |
|
return self.crit(self.glb[time, :], self.constrains[time, :]) |
|
|
|
def all_loss(self): |
|
res = [self.tloss(t).detach().numpy() for t in range(self.constrains.shape[0])] |
|
return np.array(res) |
|
|
|
''' |
|
rotation should have shape batch_size * Joint_num * (3/4) * Time |
|
position should have shape batch_size * 3 * Time |
|
offset should have shape batch_size * Joint_num * 3 |
|
output have shape batch_size * Time * Joint_num * 3 |
|
''' |
|
|
|
def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch.Tensor, order='xyz', quater=False, |
|
world=True): |
|
''' |
|
if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation') |
|
if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation') |
|
rotation = rotation.permute(0, 3, 1, 2) |
|
position = position.permute(0, 2, 1) |
|
''' |
|
result = torch.empty(rotation.shape[:-1] + (3,), device=position.device) |
|
|
|
norm = torch.norm(rotation, dim=-1, keepdim=True) |
|
rotation = rotation / norm |
|
|
|
|
|
transform = self.transform_from_quaternion(rotation) |
|
|
|
|
|
|
|
offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1)) |
|
|
|
result[..., 0, :] = position |
|
for i, pi in enumerate(self.parents): |
|
if pi == -1: |
|
assert i == 0 |
|
continue |
|
|
|
result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze() |
|
transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone()) |
|
if world: result[..., i, :] += result[..., pi, :] |
|
return result |
|
|
|
@staticmethod |
|
def transform_from_axis(euler, axis): |
|
transform = torch.empty(euler.shape[0:3] + (3, 3), device=euler.device) |
|
cos = torch.cos(euler) |
|
sin = torch.sin(euler) |
|
cord = ord(axis) - ord('x') |
|
|
|
transform[..., cord, :] = transform[..., :, cord] = 0 |
|
transform[..., cord, cord] = 1 |
|
|
|
if axis == 'x': |
|
transform[..., 1, 1] = transform[..., 2, 2] = cos |
|
transform[..., 1, 2] = -sin |
|
transform[..., 2, 1] = sin |
|
if axis == 'y': |
|
transform[..., 0, 0] = transform[..., 2, 2] = cos |
|
transform[..., 0, 2] = sin |
|
transform[..., 2, 0] = -sin |
|
if axis == 'z': |
|
transform[..., 0, 0] = transform[..., 1, 1] = cos |
|
transform[..., 0, 1] = -sin |
|
transform[..., 1, 0] = sin |
|
|
|
return transform |
|
|
|
@staticmethod |
|
def transform_from_quaternion(quater: torch.Tensor): |
|
qw = quater[..., 0] |
|
qx = quater[..., 1] |
|
qy = quater[..., 2] |
|
qz = quater[..., 3] |
|
|
|
x2 = qx + qx |
|
y2 = qy + qy |
|
z2 = qz + qz |
|
xx = qx * x2 |
|
yy = qy * y2 |
|
wx = qw * x2 |
|
xy = qx * y2 |
|
yz = qy * z2 |
|
wy = qw * y2 |
|
xz = qx * z2 |
|
zz = qz * z2 |
|
wz = qw * z2 |
|
|
|
m = torch.empty(quater.shape[:-1] + (3, 3), device=quater.device) |
|
m[..., 0, 0] = 1.0 - (yy + zz) |
|
m[..., 0, 1] = xy - wz |
|
m[..., 0, 2] = xz + wy |
|
m[..., 1, 0] = xy + wz |
|
m[..., 1, 1] = 1.0 - (xx + zz) |
|
m[..., 1, 2] = yz - wx |
|
m[..., 2, 0] = xz - wy |
|
m[..., 2, 1] = yz + wx |
|
m[..., 2, 2] = 1.0 - (xx + yy) |
|
|
|
return m |
