| """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) |
| if format.lower() == "mlx": |
| return np.transpose(vol, (0, 2, 3, 1)).astype(np.float32) |
| return np.transpose(vol, (1, 0, 2, 3)).astype(np.float32) |
|
|
