| 1 | """Small, standalone WASD/IJKL camera DSL used by UCPE inference.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import math |
| 6 | |
| 7 | import numpy as np |
| 8 | import torch |
| 9 | |
| 10 | DEFAULT_TRANSLATION_SPEED = 0.025 |
| 11 | DEFAULT_ROTATION_SPEED_DEG = 0.6 |
| 12 | DEFAULT_PITCH_SPEED_DEG = 0.2 |
| 13 | DEFAULT_PITCH_LIMIT_DEG = 60.0 |
| 14 | ALLOWED_ACTION_KEYS = frozenset("wsadikjl") |
| 15 | |
| 16 | |
| 17 | def parse_action_string(action: str) -> list[list[str]]: |
| 18 | cleaned = "".join(action.replace("，", ",").split()) |
| 19 | if not cleaned: |
| 20 | raise ValueError("action string is empty") |
| 21 | frames: list[list[str]] = [] |
| 22 | for segment in cleaned.split(","): |
| 23 | if "-" not in segment: |
| 24 | raise ValueError(f"Invalid action segment {segment!r}; expected '<keys>-<duration>'") |
| 25 | keys_part, duration = segment.rsplit("-", 1) |
| 26 | if not duration.isdigit() or int(duration) <= 0: |
| 27 | raise ValueError(f"Invalid action duration in {segment!r}") |
| 28 | keys = [] if keys_part.lower() == "none" else sorted(set(keys_part.lower())) |
| 29 | invalid = sorted(set(keys) - ALLOWED_ACTION_KEYS) |
| 30 | if invalid: |
| 31 | raise ValueError(f"Unknown action keys {invalid}; allowed: {''.join(sorted(ALLOWED_ACTION_KEYS))}") |
| 32 | frames.extend([keys] * int(duration)) |
| 33 | return frames |
| 34 | |
| 35 | |
| 36 | def _rot_x(angle: float) -> np.ndarray: |
| 37 | c, s = math.cos(angle), math.sin(angle) |
| 38 | return np.array([[1, 0, 0], [0, c, -s], [0, s, c]], dtype=np.float64) |
| 39 | |
| 40 | |
| 41 | def _rot_y(angle: float) -> np.ndarray: |
| 42 | c, s = math.cos(angle), math.sin(angle) |
| 43 | return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=np.float64) |
| 44 | |
| 45 | |
| 46 | def action_string_to_c2w( |
| 47 | frames: list[list[str]], *, translation_speed: float, rotation_speed_deg: float, |
| 48 | pitch_speed_deg: float, pitch_limit_deg: float, fps: float, |
| 49 | ) -> np.ndarray: |
| 50 | pose = np.eye(4, dtype=np.float64) |
| 51 | pitch = 0.0 |
| 52 | velocity = np.zeros(4, dtype=np.float64) |
| 53 | poses = [pose.copy()] |
| 54 | previous: set[str] = set() |
| 55 | dt = 1.0 / fps |
| 56 | for keys in frames: |
| 57 | current = set(keys) |
| 58 | target = np.array([ |
| 59 | float("w" in current) - float("s" in current), |
| 60 | float("d" in current) - float("a" in current), |
| 61 | float("l" in current) - float("j" in current), |
| 62 | float("i" in current) - float("k" in current), |
| 63 | ]) |
| 64 | target *= np.array([translation_speed, translation_speed, |
| 65 | math.radians(rotation_speed_deg), math.radians(pitch_speed_deg)]) |
| 66 | if current - previous: |
| 67 | velocity = target |
| 68 | else: |
| 69 | velocity += (target - velocity) * (1.0 - math.exp(-dt / (0.45 if np.any(target) else 1.0))) |
| 70 | previous = current |
| 71 | new_pitch = np.clip(pitch + velocity[3], -math.radians(pitch_limit_deg), math.radians(pitch_limit_deg)) |
| 72 | pitch_step = new_pitch - pitch |
| 73 | pitch = new_pitch |
| 74 | rotation = _rot_y(velocity[2]) @ pose[:3, :3] @ _rot_x(pitch_step) |
| 75 | forward = rotation[:, 2].copy(); forward[1] = 0 |
| 76 | right = rotation[:, 0].copy(); right[1] = 0 |
| 77 | forward /= max(np.linalg.norm(forward), 1e-6) |
| 78 | right /= max(np.linalg.norm(right), 1e-6) |
| 79 | pose = np.eye(4, dtype=np.float64) |
| 80 | pose[:3, :3] = rotation |
| 81 | pose[:3, 3] = poses[-1][:3, 3] + forward * velocity[0] + right * velocity[1] |
| 82 | poses.append(pose.copy()) |
| 83 | return np.stack(poses).astype(np.float32) |
| 84 | |
| 85 | |
| 86 | def default_k_pix(width: int, height: int, fov_deg: float) -> torch.Tensor: |
| 87 | focal = (width / 2.0) / math.tan(math.radians(fov_deg) / 2.0) |
| 88 | return torch.tensor([[focal, 0.0, width / 2.0], [0.0, focal, height / 2.0], [0.0, 0.0, 1.0]]) |
| 89 | |
| 90 | |
| 91 | def build_action_pt_from_string( |
| 92 | action: str, *, num_frames: int, image_width: int, image_height: int, |
| 93 | translation_speed: float, rotation_speed_deg: float, pitch_limit_deg: float, |
| 94 | fov_deg: float, fps: float, |
| 95 | ) -> dict[str, torch.Tensor | str]: |
| 96 | keys = parse_action_string(action) |
| 97 | c2ws = action_string_to_c2w( |
| 98 | keys, translation_speed=translation_speed, rotation_speed_deg=rotation_speed_deg, |
| 99 | pitch_speed_deg=DEFAULT_PITCH_SPEED_DEG, pitch_limit_deg=pitch_limit_deg, fps=fps, |
| 100 | ) |
| 101 | if len(c2ws) < num_frames: |
| 102 | c2ws = np.concatenate([c2ws, np.repeat(c2ws[-1:], num_frames - len(c2ws), axis=0)]) |
| 103 | c2ws = c2ws[:num_frames] |
| 104 | return {"c2ws_raw": torch.from_numpy(c2ws), "K_pix": default_k_pix(image_width, image_height, fov_deg), |
| 105 | "schema": "wasd_ijkl_ucpe_v3"} |
| 106 |