364 lines
14 KiB
Python
364 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from fastvideo.models.dits.lingbotworld.cam_utils import (
|
|
compute_relative_poses,
|
|
get_plucker_embeddings,
|
|
interpolate_camera_poses,
|
|
)
|
|
|
|
WSAD_OFFSET = 12.35
|
|
DIAGONAL_OFFSET = 8.73
|
|
MOUSE_PITCH_SENSITIVITY = 15.0
|
|
MOUSE_YAW_SENSITIVITY = 15.0
|
|
MOUSE_THRESHOLD = 0.02
|
|
|
|
|
|
def compute_next_pose_from_action(current_pose, keyboard_action, mouse_action):
|
|
x, y, z, pitch, yaw = current_pose
|
|
w, s, a, d = keyboard_action[:4]
|
|
mouse_x, mouse_y = mouse_action[:2]
|
|
|
|
delta_pitch = MOUSE_PITCH_SENSITIVITY * mouse_x if abs(mouse_x) >= MOUSE_THRESHOLD else 0.0
|
|
delta_yaw = MOUSE_YAW_SENSITIVITY * mouse_y if abs(mouse_y) >= MOUSE_THRESHOLD else 0.0
|
|
new_pitch = pitch + delta_pitch
|
|
new_yaw = yaw + delta_yaw
|
|
|
|
new_yaw = (new_yaw + 180) % 360 - 180
|
|
|
|
local_forward = 0.0
|
|
if w > 0.5 and s < 0.5:
|
|
local_forward = WSAD_OFFSET
|
|
elif s > 0.5 and w < 0.5:
|
|
local_forward = -WSAD_OFFSET
|
|
|
|
local_right = 0.0
|
|
if d > 0.5 and a < 0.5:
|
|
local_right = WSAD_OFFSET
|
|
elif a > 0.5 and d < 0.5:
|
|
local_right = -WSAD_OFFSET
|
|
|
|
if abs(local_forward) > 0.1 and abs(local_right) > 0.1:
|
|
local_forward = np.sign(local_forward) * DIAGONAL_OFFSET
|
|
local_right = np.sign(local_right) * DIAGONAL_OFFSET
|
|
|
|
avg_yaw = float((yaw + new_yaw) / 2.0)
|
|
yaw_rad = float(np.deg2rad(avg_yaw))
|
|
cos_yaw = np.cos(yaw_rad)
|
|
sin_yaw = np.sin(yaw_rad)
|
|
|
|
delta_x = cos_yaw * local_forward - sin_yaw * local_right
|
|
delta_y = sin_yaw * local_forward + cos_yaw * local_right
|
|
return np.array([x + delta_x, y + delta_y, z, new_pitch, new_yaw], dtype=np.float32)
|
|
|
|
|
|
def compute_all_poses_from_actions(keyboard_conditions, mouse_conditions):
|
|
poses = np.zeros((len(keyboard_conditions), 5), dtype=np.float32)
|
|
for idx in range(len(keyboard_conditions) - 1):
|
|
poses[idx + 1] = compute_next_pose_from_action(poses[idx], keyboard_conditions[idx], mouse_conditions[idx])
|
|
return poses
|
|
|
|
|
|
def build_intrinsics(height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
|
fov_rad = float(np.deg2rad(90.0))
|
|
fx = float(width) / (2.0 * float(np.tan(fov_rad / 2.0)))
|
|
fy = float(height) / (2.0 * float(np.tan(fov_rad / 2.0)))
|
|
return torch.tensor([[fx, fy, width / 2.0, height / 2.0]], device=device, dtype=dtype)
|
|
|
|
|
|
def build_matrixgame3_action_preset(num_frames: int, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
|
|
presets = create_action_presets(num_frames=num_frames, keyboard_dim=6, seed=seed)
|
|
return presets["keyboard"], presets["mouse"]
|
|
|
|
|
|
def interpolate_camera_poses_handedness(
|
|
src_indices: np.ndarray,
|
|
src_rot_mat: np.ndarray,
|
|
src_trans_vec: np.ndarray,
|
|
tgt_indices: np.ndarray,
|
|
) -> torch.Tensor:
|
|
dets = np.linalg.det(src_rot_mat)
|
|
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
|
|
if flip_handedness:
|
|
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
|
|
src_rot_mat = src_rot_mat @ flip_mat
|
|
|
|
c2ws = interpolate_camera_poses(
|
|
src_indices=src_indices,
|
|
src_rot_mat=src_rot_mat,
|
|
src_trans_vec=src_trans_vec,
|
|
tgt_indices=tgt_indices,
|
|
)
|
|
if flip_handedness:
|
|
flip_mat_t = torch.from_numpy(flip_mat).to(c2ws.device, dtype=c2ws.dtype)
|
|
c2ws[:, :3, :3] = c2ws[:, :3, :3] @ flip_mat_t
|
|
return c2ws
|
|
|
|
|
|
def build_extrinsics_from_actions(
|
|
keyboard_conditions: torch.Tensor,
|
|
mouse_conditions: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
keyboard_np = keyboard_conditions.detach().to(dtype=torch.float32).cpu().numpy()
|
|
mouse_np = mouse_conditions.detach().to(dtype=torch.float32).cpu().numpy()
|
|
poses = compute_all_poses_from_actions(keyboard_np, mouse_np)
|
|
rotations = np.concatenate(
|
|
[np.zeros((poses.shape[0], 1), dtype=np.float32), poses[:, 3:5]],
|
|
axis=1,
|
|
)
|
|
positions = poses[:, :3]
|
|
return build_extrinsics(rotations, positions)
|
|
|
|
|
|
def build_extrinsics(
|
|
video_rotation: np.ndarray,
|
|
video_position: np.ndarray,
|
|
) -> torch.Tensor:
|
|
extrinsics = []
|
|
for frame_rotation, frame_position in zip(video_rotation, video_position, strict=False):
|
|
roll_deg, pitch_deg, yaw_deg = frame_rotation
|
|
roll, pitch, yaw = np.radians([roll_deg, pitch_deg, yaw_deg])
|
|
|
|
rot_z = np.array(
|
|
[[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]],
|
|
dtype=np.float32,
|
|
)
|
|
rot_y = np.array(
|
|
[[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]],
|
|
dtype=np.float32,
|
|
)
|
|
rot_x = np.array(
|
|
[[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]],
|
|
dtype=np.float32,
|
|
)
|
|
rot = rot_z @ rot_y @ rot_x
|
|
ext = np.eye(4, dtype=np.float32)
|
|
ext[:3, :3] = rot
|
|
ext[:3, 3] = np.asarray(frame_position, dtype=np.float32)
|
|
extrinsics.append(ext)
|
|
|
|
r_init = np.array([[0, 0, 1], [1, 0, 0], [0, -1, 0]], dtype=np.float32)
|
|
extrinsics_tensor = torch.from_numpy(np.stack(extrinsics, axis=0)).float()
|
|
extrinsics_tensor[:, :3, :3] = extrinsics_tensor[:, :3, :3] @ torch.from_numpy(r_init)
|
|
extrinsics_tensor[:, :3, 3] = extrinsics_tensor[:, :3, 3] * 0.01
|
|
return extrinsics_tensor
|
|
|
|
|
|
def build_plucker_from_c2ws(
|
|
c2ws_seq: torch.Tensor,
|
|
src_indices: np.ndarray,
|
|
tgt_indices: np.ndarray,
|
|
*,
|
|
target_h: int,
|
|
target_w: int,
|
|
latent_h: int,
|
|
latent_w: int,
|
|
framewise: bool = True,
|
|
) -> torch.Tensor:
|
|
c2ws_np = c2ws_seq.detach().to(dtype=torch.float32).cpu().numpy()
|
|
c2ws_infer = interpolate_camera_poses_handedness(
|
|
src_indices=src_indices,
|
|
src_rot_mat=c2ws_np[:, :3, :3],
|
|
src_trans_vec=c2ws_np[:, :3, 3],
|
|
tgt_indices=tgt_indices,
|
|
).to(device=c2ws_seq.device, dtype=c2ws_seq.dtype)
|
|
c2ws_infer = compute_relative_poses(c2ws_infer, framewise=framewise)
|
|
return build_plucker_from_pose(
|
|
c2ws_infer,
|
|
target_h=target_h,
|
|
target_w=target_w,
|
|
latent_h=latent_h,
|
|
latent_w=latent_w,
|
|
)
|
|
|
|
|
|
def build_plucker_from_pose(
|
|
c2ws_pose: torch.Tensor,
|
|
*,
|
|
target_h: int,
|
|
target_w: int,
|
|
latent_h: int,
|
|
latent_w: int,
|
|
) -> torch.Tensor:
|
|
ks = build_intrinsics(target_h, target_w, c2ws_pose.device, c2ws_pose.dtype).repeat(c2ws_pose.shape[0], 1)
|
|
plucker = get_plucker_embeddings(c2ws_pose, ks, target_h, target_w)
|
|
c1 = target_h // latent_h
|
|
c2 = target_w // latent_w
|
|
plucker = plucker.view(c2ws_pose.shape[0], latent_h, c1, latent_w, c2, 6)
|
|
plucker = plucker.permute(0, 1, 3, 5, 2, 4).reshape(c2ws_pose.shape[0], latent_h, latent_w, 6 * c1 * c2)
|
|
plucker = plucker.permute(3, 0, 1, 2).unsqueeze(0)
|
|
return plucker
|
|
|
|
|
|
def select_memory_idx_fov(
|
|
extrinsics_all: torch.Tensor | np.ndarray,
|
|
current_start_frame_idx: int,
|
|
selected_index_base: list[int],
|
|
*,
|
|
height: int = 720,
|
|
width: int = 1280,
|
|
return_confidence: bool = False,
|
|
) -> list[int] | tuple[list[int], list[float]]:
|
|
if isinstance(extrinsics_all, np.ndarray):
|
|
extrinsics_tensor = torch.from_numpy(extrinsics_all).float()
|
|
else:
|
|
extrinsics_tensor = extrinsics_all.float()
|
|
|
|
device = extrinsics_tensor.device
|
|
if current_start_frame_idx <= 1:
|
|
selected = [0] * len(selected_index_base)
|
|
if return_confidence:
|
|
return selected, [0.0] * len(selected)
|
|
return selected
|
|
|
|
fov_rad = np.deg2rad(90.0)
|
|
fx = width / (2 * np.tan(fov_rad / 2))
|
|
fy = height / (2 * np.tan(fov_rad / 2))
|
|
near, far = 0.1, 30.0
|
|
|
|
candidate_indices = torch.arange(1, current_start_frame_idx, device=device)
|
|
r_cand = extrinsics_tensor[candidate_indices, :3, :3]
|
|
t_cand = extrinsics_tensor[candidate_indices, :3, 3:4]
|
|
r_cand_inv = r_cand.transpose(1, 2)
|
|
t_cand_inv = -torch.bmm(r_cand_inv, t_cand)
|
|
|
|
num_side = 10
|
|
z_samples = torch.linspace(near, far, num_side, device=device)
|
|
x_samples = torch.linspace(-1, 1, num_side, device=device)
|
|
y_samples = torch.linspace(-1, 1, num_side, device=device)
|
|
grid_x, grid_y, grid_z = torch.meshgrid(x_samples, y_samples, z_samples, indexing="ij")
|
|
points_cam_base = torch.stack([
|
|
grid_x.reshape(-1) * grid_z.reshape(-1) * (width / (2 * fx)),
|
|
grid_y.reshape(-1) * grid_z.reshape(-1) * (height / (2 * fy)),
|
|
grid_z.reshape(-1),
|
|
],
|
|
dim=0)
|
|
|
|
selected_index: list[int] = []
|
|
selected_confidence: list[float] = []
|
|
for frame_idx in selected_index_base:
|
|
base_pose = extrinsics_tensor[frame_idx]
|
|
points_world = base_pose[:3, :3] @ points_cam_base + base_pose[:3, 3:4]
|
|
points_world_batched = points_world.unsqueeze(0)
|
|
points_in_cands = torch.bmm(r_cand_inv, points_world_batched.expand(len(candidate_indices), -1,
|
|
-1)) + t_cand_inv
|
|
|
|
x = points_in_cands[:, 0, :]
|
|
y = points_in_cands[:, 1, :]
|
|
z = points_in_cands[:, 2, :]
|
|
u = (x * fx / torch.clamp(z, min=1e-6)) + width / 2
|
|
v = (y * fy / torch.clamp(z, min=1e-6)) + height / 2
|
|
|
|
in_view = (z > near) & (z < far) & (u >= 0) & (u <= width) & (v >= 0) & (v <= height)
|
|
ratios = in_view.float().mean(dim=1)
|
|
best_idx = torch.argmax(ratios)
|
|
selected_index.append(candidate_indices[best_idx].item())
|
|
selected_confidence.append(ratios[best_idx].item())
|
|
|
|
if return_confidence:
|
|
return selected_index, selected_confidence
|
|
return selected_index
|
|
|
|
|
|
def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None):
|
|
if keyboard_dim not in (2, 4, 6, 7):
|
|
raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}")
|
|
if num_frames % 4 != 1:
|
|
raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.")
|
|
|
|
if seed is not None:
|
|
random.seed(seed)
|
|
|
|
num_samples_per_action = 4
|
|
|
|
if keyboard_dim == 4:
|
|
actions_single_action = ["forward", "left", "right"]
|
|
actions_double_action = ["forward_left", "forward_right"]
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3}
|
|
elif keyboard_dim == 2:
|
|
actions_single_action = ["forward", "back"]
|
|
actions_double_action = []
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1}
|
|
elif keyboard_dim == 6:
|
|
actions_single_action = ["forward", "back", "left", "right"]
|
|
actions_double_action = ["forward_left", "forward_right"]
|
|
actions_single_camera = ["camera_l", "camera_r"]
|
|
keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5}
|
|
else: # keyboard_dim == 7
|
|
actions_single_action = ["forward", "back", "left", "right"]
|
|
actions_double_action = []
|
|
actions_single_camera = []
|
|
keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6}
|
|
|
|
actions_to_test = (actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5)
|
|
for action in (actions_single_action + actions_double_action):
|
|
for camera in actions_single_camera:
|
|
actions_to_test.append(f"{action}_{camera}")
|
|
|
|
if not actions_to_test:
|
|
actions_to_test = actions_single_action * 5
|
|
|
|
base_action = actions_single_action + actions_single_camera
|
|
cam_value = 0.1
|
|
camera_value_map = {
|
|
"camera_up": [cam_value, 0],
|
|
"camera_down": [-cam_value, 0],
|
|
"camera_l": [0, -cam_value],
|
|
"camera_r": [0, cam_value],
|
|
"camera_ur": [cam_value, cam_value],
|
|
"camera_ul": [cam_value, -cam_value],
|
|
"camera_dr": [-cam_value, cam_value],
|
|
"camera_dl": [-cam_value, -cam_value],
|
|
}
|
|
|
|
data = []
|
|
for action_name in actions_to_test:
|
|
keyboard_condition = torch.zeros((num_samples_per_action, keyboard_dim))
|
|
mouse_condition = torch.zeros((num_samples_per_action, 2))
|
|
|
|
for sub_act in base_action:
|
|
if sub_act not in action_name:
|
|
continue
|
|
if sub_act in camera_value_map:
|
|
mouse_condition = torch.tensor(
|
|
[camera_value_map[sub_act] for _ in range(num_samples_per_action)],
|
|
dtype=mouse_condition.dtype,
|
|
)
|
|
elif sub_act in keyboard_idx:
|
|
keyboard_condition[:, keyboard_idx[sub_act]] = 1
|
|
|
|
data.append({
|
|
"keyboard_condition": keyboard_condition,
|
|
"mouse_condition": mouse_condition,
|
|
})
|
|
|
|
keyboard_condition = torch.zeros((num_frames, keyboard_dim))
|
|
mouse_condition = torch.zeros((num_frames, 2))
|
|
current_frame = 0
|
|
selections = [12]
|
|
|
|
while current_frame < num_frames:
|
|
rd_frame = selections[random.randint(0, len(selections) - 1)]
|
|
entry = data[random.randint(0, len(data) - 1)]
|
|
key_seq = entry["keyboard_condition"]
|
|
mouse_seq = entry["mouse_condition"]
|
|
|
|
if current_frame == 0:
|
|
keyboard_condition[:1] = key_seq[:1]
|
|
mouse_condition[:1] = mouse_seq[:1]
|
|
current_frame = 1
|
|
else:
|
|
rd_frame = min(rd_frame, num_frames - current_frame)
|
|
repeat_time = rd_frame // 4
|
|
keyboard_condition[current_frame:current_frame + rd_frame] = key_seq.repeat(repeat_time, 1)
|
|
mouse_condition[current_frame:current_frame + rd_frame] = mouse_seq.repeat(repeat_time, 1)
|
|
current_frame += rd_frame
|
|
|
|
return {"keyboard": keyboard_condition, "mouse": mouse_condition}
|