Spaces:
Running on Zero
Running on Zero
File size: 1,997 Bytes
1d998f6 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import torch
from .rotation import mat_to_quat, quat_to_mat
def extri_intri_to_pose_encoding(
extrinsics,
intrinsics,
image_size_hw=None,
pose_encoding_type="absT_quaR_FoV",
):
if pose_encoding_type != "absT_quaR_FoV":
raise NotImplementedError
rotation = extrinsics[:, :, :3, :3]
translation = extrinsics[:, :, :3, 3]
quat = mat_to_quat(rotation)
height, width = image_size_hw
fov_h = 2 * torch.atan((height / 2) / intrinsics[..., 1, 1])
fov_w = 2 * torch.atan((width / 2) / intrinsics[..., 0, 0])
return torch.cat(
[translation, quat, fov_h[..., None], fov_w[..., None]],
dim=-1,
).float()
def _pose_encoding_to_extri_intri(
pose_encoding,
image_size_hw=None,
pose_encoding_type="absT_quaR_FoV",
build_intrinsics=True,
):
if pose_encoding_type != "absT_quaR_FoV":
raise NotImplementedError
translation = pose_encoding[..., :3]
quat = pose_encoding[..., 3:7]
fov_h = pose_encoding[..., 7]
fov_w = pose_encoding[..., 8]
rotation = quat_to_mat(quat)
extrinsics = torch.cat([rotation, translation[..., None]], dim=-1)
intrinsics = None
if build_intrinsics:
height, width = image_size_hw
fy = (height / 2.0) / torch.tan(fov_h / 2.0)
fx = (width / 2.0) / torch.tan(fov_w / 2.0)
intrinsics = torch.zeros(
pose_encoding.shape[:2] + (3, 3),
device=pose_encoding.device,
)
intrinsics[..., 0, 0] = fx
intrinsics[..., 1, 1] = fy
intrinsics[..., 0, 2] = width / 2
intrinsics[..., 1, 2] = height / 2
intrinsics[..., 2, 2] = 1.0
return extrinsics, intrinsics
def pose_encoding_to_extri_intri(
pose_encoding,
image_size_hw=None,
pose_encoding_type="absT_quaR_FoV",
build_intrinsics=True,
):
return _pose_encoding_to_extri_intri(
pose_encoding,
image_size_hw,
pose_encoding_type,
build_intrinsics,
)
|