Spaces:
Running on Zero
Running on Zero
File size: 5,865 Bytes
c1e2af3 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Forward-kinematics primitives for articulated skeletons."""
from typing import List
import einops
import torch
import torch.nn.functional as F
from ..tools import ensure_batched
@ensure_batched(local_joint_rots=4, root_positions=2)
def fk(
local_joint_rots: torch.Tensor,
root_positions: torch.Tensor,
skeleton,
root_positions_is_global: bool = True,
):
"""Compute global joint rotations and positions from local rotations.
Args:
local_joint_rots: Local rotation matrices with shape `(..., J, 3, 3)`.
root_positions: Root translations with shape `(..., 3)`.
skeleton: Skeleton object exposing `neutral_joints`, `joint_parents`, and
`root_idx`.
root_positions_is_global: If `True`, neutral joints are recentered so root
translations are interpreted in world space.
Returns:
Tuple `(global_joint_rots, posed_joints, posed_joints_norootpos)`.
"""
device = local_joint_rots.device
dtype = local_joint_rots.dtype
neutral_joints = skeleton.neutral_joints.to(device=device, dtype=dtype)
if root_positions_is_global is True:
# Removing the pelvis offset from the neutral joints
# as the root positions does not depends on the pelvis offset of the skeleton
pelvis_offset = neutral_joints[skeleton.root_idx]
neutral_joints = neutral_joints - pelvis_offset
# compute joint position and global rotations
joints = einops.repeat(
neutral_joints,
"j k -> b j k",
b=len(local_joint_rots),
)
posed_joints_norootpos, global_joint_rots = batch_rigid_transform(
local_joint_rots,
joints,
skeleton.joint_parents,
skeleton.root_idx,
)
# if root_positions_is_global is True:
# posed_joints_norootpos always start at zero
# otherwise it could start with the pelvis offset
posed_joints = posed_joints_norootpos + root_positions[:, None]
return global_joint_rots, posed_joints, posed_joints_norootpos
def compute_idx_levels(parents):
"""Group joint indices by hierarchy depth for level-wise FK updates.
Args:
parents: Parent index tensor of shape `(J,)` with root parent `-1`.
Returns:
List of index tensors, where each tensor contains joints at one depth.
"""
idx_levs = [[]]
lev_dicts = {0: -1}
for i in range(1, parents.shape[0]):
assert int(parents[i]) in lev_dicts
lev = lev_dicts[int(parents[i])] + 1
if lev + 1 > len(idx_levs):
idx_levs.append([])
idx_levs[lev].append(int(i))
lev_dicts[int(i)] = lev
idx_levs = [torch.tensor(x).long() for x in idx_levs]
return idx_levs
def batch_rigid_transform(rot_mats, joints, parents, root_idx):
"""Perform batch rigid transformation on a skeletal structure.
Args:
rot_mats: Local rotation matrices for each joint: (B, J, 3, 3)
joints: Initial joint positions: (B, J, 3)
parents: Tensor indicating the parent of each joint: (J,)
root_idx (int): index of the root
Returns:
Transformed joint positions after applying forward kinematics.
"""
# Compute the hierarchical levels of joints based on their parent relationships
idx_levs = compute_idx_levels(parents)
# Apply forward kinematics to transform the joints
return forward_kinematics(rot_mats, joints, parents, idx_levs, root_idx)
@torch.jit.script
def transform_mat(R, t):
"""Creates a batch of transformation matrices.
Args:
- R: Bx3x3 array of a batch of rotation matrices
- t: Bx3x1 array of a batch of translation vectors
Returns:
- T: Bx4x4 Transformation matrix
"""
# No padding left or right, only add an extra row
return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1.0)], dim=2)
@torch.jit.script
def forward_kinematics(
rot_mats,
joints,
parents: torch.Tensor,
idx_levs: List[torch.Tensor],
root_idx: int,
):
"""Perform forward kinematics to compute posed joints and global rotation matrices.
Args:
rot_mats: Local rotation matrices for each joint: (B, J, 3, 3)
joints: Initial joint positions: (B, J, 3)
parents: Tensor indicating the parent of each joint: (J,)
idx_levs: Tensors of joint indices grouped by depth in the kinematic tree.
root_idx (int): index of the root
Returns:
Posed joints: (B, J, 3)
Global rotation matrices: (B, J, 3, 3)
"""
# Add an extra dimension to joints
joints = torch.unsqueeze(joints, dim=-1)
# Compute relative joint positions
rel_joints = joints.clone()
mask_no_root = torch.ones(joints.shape[1], dtype=torch.bool)
mask_no_root[root_idx] = False
rel_joints[:, mask_no_root] -= joints[:, parents[mask_no_root]].clone()
# Compute initial transformation matrices
# (B, J + 1, 4, 4)
transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3), rel_joints.reshape(-1, 3, 1)).reshape(
-1, joints.shape[1], 4, 4
)
# Initialize the root transformation matrices
transforms = torch.zeros_like(transforms_mat)
transforms[:, root_idx] = transforms_mat[:, root_idx]
# Compute global transformations level by level
for indices in idx_levs:
curr_res = torch.matmul(transforms[:, parents[indices]], transforms_mat[:, indices])
transforms[:, indices] = curr_res
# Extract posed joint positions from the transformation matrices
posed_joints = transforms[:, :, :3, 3]
# Extract global rotation matrices from the transformation matrices
global_rot_mat = transforms[:, :, :3, :3]
return posed_joints, global_rot_mat
|