Spaces:
Sleeping
Sleeping
| import torch | |
| from longstream.utils.vendor.models.components.utils.rotation import ( | |
| quat_to_mat, | |
| mat_to_quat, | |
| ) | |
| def compose_abs_from_rel( | |
| rel_pose_enc: torch.Tensor, keyframe_indices: torch.Tensor | |
| ) -> torch.Tensor: | |
| squeeze_batch = False | |
| if rel_pose_enc.ndim == 2: | |
| rel_pose_enc = rel_pose_enc.unsqueeze(0) | |
| squeeze_batch = True | |
| if keyframe_indices.ndim == 1: | |
| keyframe_indices = keyframe_indices.unsqueeze(0) | |
| if rel_pose_enc.ndim != 3 or keyframe_indices.ndim != 2: | |
| raise ValueError( | |
| f"Expected rel_pose_enc [B,S,D] or [S,D] and keyframe_indices [B,S] or [S], " | |
| f"got {tuple(rel_pose_enc.shape)} and {tuple(keyframe_indices.shape)}" | |
| ) | |
| B, S, _ = rel_pose_enc.shape | |
| device = rel_pose_enc.device | |
| dtype = rel_pose_enc.dtype | |
| rel_t = rel_pose_enc[..., :3] | |
| rel_q = rel_pose_enc[..., 3:7] | |
| rel_f = rel_pose_enc[..., 7:9] | |
| rel_R = quat_to_mat(rel_q.reshape(-1, 4)).reshape(B, S, 3, 3) | |
| abs_R = torch.zeros(B, S, 3, 3, device=device, dtype=dtype) | |
| abs_t = torch.zeros(B, S, 3, device=device, dtype=dtype) | |
| abs_f = torch.zeros(B, S, 2, device=device, dtype=dtype) | |
| for b in range(B): | |
| abs_R[b, 0] = rel_R[b, 0] | |
| abs_t[b, 0] = rel_t[b, 0] | |
| abs_f[b, 0] = rel_f[b, 0] | |
| for s in range(1, S): | |
| ref_idx = int(keyframe_indices[b, s].item()) | |
| abs_R[b, s] = rel_R[b, s] @ abs_R[b, ref_idx] | |
| abs_t[b, s] = rel_t[b, s] + rel_R[b, s] @ abs_t[b, ref_idx] | |
| abs_f[b, s] = rel_f[b, s] | |
| abs_q = mat_to_quat(abs_R.reshape(-1, 3, 3)).reshape(B, S, 4) | |
| abs_pose_enc = torch.cat([abs_t, abs_q, abs_f], dim=-1) | |
| if squeeze_batch: | |
| return abs_pose_enc[0] | |
| return abs_pose_enc | |