ZhenweiWang's picture
Upload folder using huggingface_hub
0ca05b5 verified
import torch
from .rotation import quat_to_rotmat, rotmat_to_quat
def camera_params_to_vector(
ext, intr, image_hw=None
):
"""Convert camera matrices to a compact vector."""
# ext: (..., 3, 4): Camera-to-world extrinsic [R|t]
# intr: (..., 3, 3): Intrinsics
# image_hw: (h, w)
R = ext[..., :3, :3] # Rotation part
t = ext[..., :3, 3] # Translation part
q = rotmat_to_quat(R) # Quaternion (wxyz)
h, w = image_hw
fov_v = 2.0 * torch.atan(h * 0.5 / intr[..., 1, 1]) # Vertical FOV
fov_u = 2.0 * torch.atan(w * 0.5 / intr[..., 0, 0]) # Horizontal FOV
vec = torch.stack([
t[..., 0], t[..., 1], t[..., 2],
q[..., 0], q[..., 1], q[..., 2], q[..., 3],
fov_v, fov_u
], dim=-1).float()
return vec
def extrinsics_to_vector(ext):
"""Convert extrinsics to [t, q] vector."""
# ext: (..., 3, 4)
R = ext[..., :3, :3]
t = ext[..., :3, 3]
q = rotmat_to_quat(R)
vec = torch.stack([
t[..., 0], t[..., 1], t[..., 2],
q[..., 0], q[..., 1], q[..., 2], q[..., 3]
], dim=-1).float()
return vec
def vector_to_extrinsics(cam_vec):
"""Convert [t, q] vector to extrinsic matrix."""
# cam_vec: (..., 7)
q = cam_vec[..., 3:7]
t = cam_vec[..., :3]
R = quat_to_rotmat(q)
ext = torch.cat([R, t.unsqueeze(-1)], dim=-1)
return ext
def vector_to_camera_matrices(
cam_vec, image_hw=None, build_intr=True
):
"""Reconstruct extrinsic and intrinsic matrix from vector."""
# cam_vec: (..., 9)
intr = None
# Decompose vector
t = cam_vec[..., 0:3]
q = cam_vec[..., 3:7]
fov_v = cam_vec[..., 7]
fov_u = cam_vec[..., 8]
# Build extrinsic: [R|t]
R = quat_to_rotmat(q)
ext = torch.cat([R, t.unsqueeze(-1)], dim=-1)
# Build intrinsic if needed
if build_intr:
h, w = image_hw
fy = h * 0.5 / torch.tan(fov_v * 0.5)
fx = w * 0.5 / torch.tan(fov_u * 0.5)
shape = cam_vec.shape[:-1] + (3, 3)
intr = torch.zeros(shape, device=cam_vec.device, dtype=cam_vec.dtype)
intr[..., 0, 0] = fx
intr[..., 1, 1] = fy
intr[..., 0, 2] = w * 0.5
intr[..., 1, 2] = h * 0.5
intr[..., 2, 2] = 1.0
return ext, intr