File size: 6,610 Bytes
87c8483 | 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 173 174 175 176 177 178 179 180 181 | """Interpolation and frame-rate rescaling for pose sequences.
Provides linear interpolation (via scipy), quaternion slerp, and
functions to up/down-sample joint pose trajectories to a target frame rate.
"""
import torch
import numpy as np
from scipy.interpolate import interp1d
from .kornia_transform import angle_axis_to_quaternion, quaternion_to_angle_axis
def interp_tensor_with_scipy(x, new_len=None, scale=None, dim=-1):
orig_len = x.shape[dim]
if new_len is None:
new_len = int(orig_len * scale)
T = orig_len
f = interp1d(
np.linspace(0, T, orig_len),
x.cpu().numpy(),
axis=dim,
assume_sorted=True,
fill_value="extrapolate",
)
x_interp = torch.from_numpy(f(np.linspace(0, T, new_len))).type_as(x)
return x_interp
def slerp(q0, q1, t):
# type: (torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
cos_half_theta = torch.sum(q0 * q1, dim=-1)
neg_mask = cos_half_theta < 0
q1 = q1.clone()
# Replace: q1[neg_mask] = -q1[neg_mask]
# With: torch.where for safer broadcasting
neg_mask_expanded = neg_mask.unsqueeze(-1).expand_as(q1)
q1 = torch.where(neg_mask_expanded, -q1, q1)
cos_half_theta = torch.abs(cos_half_theta)
cos_half_theta = torch.unsqueeze(cos_half_theta, dim=-1)
half_theta = torch.acos(cos_half_theta)
sin_half_theta = torch.sqrt(1.0 - cos_half_theta * cos_half_theta)
ratioA = torch.sin((1 - t[:, None]) * half_theta) / sin_half_theta
ratioB = torch.sin(t[:, None] * half_theta) / sin_half_theta
new_q = ratioA * q0 + ratioB * q1
new_q = torch.where(torch.abs(sin_half_theta) < 0.001, 0.5 * q0 + 0.5 * q1, new_q)
new_q = torch.where(torch.abs(cos_half_theta) >= 1, q0, new_q)
return new_q
def _slerp_batch(a: torch.Tensor, b: torch.Tensor, blend: torch.Tensor) -> torch.Tensor:
"""Spherical linear interpolation between two quaternions."""
slerped_quats = torch.zeros_like(a)
slerped_quats = slerp(a, b, blend)
return slerped_quats
def interpolate_quaternions(
pose_quat: torch.Tensor, source_fps: float, target_fps: float
) -> torch.Tensor:
"""
Interpolate quaternions from source_fps to target_fps.
Args:
pose_quat: Input quaternions, shape (1, T, 4)
source_fps: Source frame rate
target_fps: Target frame rate
Returns:
Interpolated quaternions
"""
device = pose_quat.device
in_shape = pose_quat.shape
assert in_shape[0] == 1, "Only support single sequence for now"
T = in_shape[1]
duration = (T - 1) * (1 / source_fps)
times = torch.arange(0, duration + 1e-6, 1 / target_fps, dtype=torch.float32, device=device)
times = times[times <= duration]
# Compute frame indices and blend factors
frame_indices = times * source_fps
index_0 = torch.floor(frame_indices).long()
index_1 = torch.min(index_0 + 1, torch.tensor(T - 1, device=device))
blend = frame_indices - index_0.float()
pose_quat_interp = _slerp_batch(pose_quat[0, index_0], pose_quat[0, index_1], blend)
pose_quat_interp = pose_quat_interp.unsqueeze(0)
return pose_quat_interp
def interpolate_pose(
pose_aa: torch.Tensor,
source_fps: float,
target_fps: float,
device: str = "cpu",
interpolation_type: str = "slerp",
rot_type: str = "aa",
) -> torch.Tensor:
"""
Interpolate pose_aa from source_fps to target_fps using specified interpolation method.
Args:
pose_aa: Input pose in angle-axis format, shape (T, N*3) where T is number of frames and N is number of joints
source_fps: Source frame rate
target_fps: Target frame rate
device: Device to run computations on
interpolation_type: Type of interpolation to use ("linear" or "slerp")
Returns:
Interpolated pose_aa with new frame rate, shape (T_new, N*3)
"""
# pose_aa: (T, N*3)
orig_shape = pose_aa.shape[1:]
if pose_aa.ndim != 2:
pose_aa = pose_aa.reshape(pose_aa.shape[0], -1)
T, D = pose_aa.shape
if interpolation_type == "linear":
# Direct linear interpolation on angle-axis representation
duration = (T - 1) * (1 / source_fps)
times = torch.arange(0, duration + 1e-6, 1 / target_fps, dtype=torch.float32, device=device)
times = times[times <= duration]
# Compute frame indices and blend factors for linear interpolation
frame_indices = times * source_fps
index_0 = torch.floor(frame_indices).long()
index_1 = torch.min(index_0 + 1, torch.tensor(T - 1, device=device))
blend = frame_indices - index_0.float()
# Linear interpolation on the entire 2D tensor
pose_aa_interp = (1 - blend.unsqueeze(1)) * pose_aa[index_0] + blend.unsqueeze(1) * pose_aa[
index_1
]
pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], *orig_shape)
if pose_aa.dtype == torch.int64:
pose_aa_interp = pose_aa_interp.round()
pose_aa_interp = pose_aa_interp.type_as(pose_aa)
return pose_aa_interp
elif interpolation_type == "slerp":
dim = 3 if rot_type == "aa" else 4
N = D // dim
pose_aa_reshaped = pose_aa.view(T, N, dim)
# Original spherical linear interpolation on quaternions
pose_aa_interp_list = []
for i in range(N):
# Convert angle-axis to quaternion for this joint
if rot_type == "aa":
pose_quat = angle_axis_to_quaternion(pose_aa_reshaped[:, i]) # (T, 4)
else:
pose_quat = pose_aa_reshaped[:, i]
pose_quat_batch = pose_quat.unsqueeze(0) # (1, T, 4)
pose_quat_interp = interpolate_quaternions(pose_quat_batch, source_fps, target_fps)
pose_quat_interp = pose_quat_interp[0] # (T_new, 4)
if rot_type == "aa":
pose_aa_interp = quaternion_to_angle_axis(pose_quat_interp) # (T_new, 3)
else:
pose_aa_interp = pose_quat_interp
pose_aa_interp_list.append(pose_aa_interp)
# Concatenate all joints: (T_new, N, 3) -> (T_new, N*3)
pose_aa_interp = torch.stack(pose_aa_interp_list, dim=1) # (T_new, N, 3)
pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], -1) # (T_new, N*3)
pose_aa_interp = pose_aa_interp.view(pose_aa_interp.shape[0], *orig_shape).to(pose_aa)
return pose_aa_interp
else:
raise ValueError(
f"Unsupported interpolation_type: {interpolation_type}. Must be 'linear' or 'slerp'."
)
|