File size: 1,808 Bytes
4fc462d | 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 | """PoseConv3D Preprocessing: (T, 17, 2) keypoints -> (1, 32, 64, 64) heatmaps."""
import numpy as np
def resample_keypoints(kp_seq: np.ndarray, target_t: int = 32) -> np.ndarray:
"""(L, J, 2) in [-1, 1] range -> (target_t, J, 2) linearly interpolated over time."""
L, J, D = kp_seq.shape
if L == target_t:
return kp_seq.copy()
src = np.arange(L)
dst = np.linspace(0, L - 1, target_t)
out = np.zeros((target_t, J, D), dtype=np.float32)
for j in range(J):
for d in range(D):
col = kp_seq[:, j, d]
bad = np.isnan(col)
if bad.all():
out[:, j, d] = np.nan
elif bad.any():
col_clean = np.interp(src, src[~bad], col[~bad])
out[:, j, d] = np.interp(dst, src, col_clean)
else:
out[:, j, d] = np.interp(dst, src, col)
return out
def rasterize_heatmaps(kp: np.ndarray, h: int = 64, w: int = 64,
sigma: float = 1.8, format: str = "pytorch") -> np.ndarray:
"""(target_t, J, 2) keypoints -> (1, target_t, h, w) for PyTorch or (target_t, h, w, 1) for MLX."""
T, J, _ = kp.shape
y_grid, x_grid = np.ogrid[:h, :w]
inv_2s2 = 1.0 / (2.0 * sigma * sigma)
cx = (kp[..., 0] + 1.0) * 0.5 * (w - 1)
cy = (kp[..., 1] + 1.0) * 0.5 * (h - 1)
dx = x_grid[None, None, :, :] - cx[:, :, None, None]
dy = y_grid[None, None, :, :] - cy[:, :, None, None]
dist_sq = dx * dx + dy * dy
g = np.exp(-dist_sq * inv_2s2)
g = np.nan_to_num(g, nan=0.0)
vol = np.max(g, axis=1, keepdims=True) # (T, 1, H, W)
if format.lower() == "mlx":
return np.transpose(vol, (0, 2, 3, 1)).astype(np.float32) # (T, H, W, 1)
return np.transpose(vol, (1, 0, 2, 3)).astype(np.float32) # (1, T, H, W)
|