Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import torch | |
| from .rotation import quat_to_mat, mat_to_quat | |
| import os | |
| import torch | |
| import numpy as np | |
| import gzip | |
| import json | |
| import random | |
| import logging | |
| import warnings | |
| from lingbot_map.utils.geometry import closed_form_inverse_se3, closed_form_inverse_se3_general | |
| def extri_intri_to_pose_encoding( | |
| extrinsics, intrinsics, image_size_hw=None, pose_encoding_type="absT_quaR_FoV" # e.g., (256, 512) | |
| ): | |
| """Convert camera extrinsics and intrinsics to a compact pose encoding. | |
| This function transforms camera parameters into a unified pose encoding format, | |
| which can be used for various downstream tasks like pose prediction or representation. | |
| Args: | |
| extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4, | |
| where B is batch size and S is sequence length. | |
| In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation. | |
| The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector. | |
| intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3. | |
| Defined in pixels, with format: | |
| [[fx, 0, cx], | |
| [0, fy, cy], | |
| [0, 0, 1]] | |
| where fx, fy are focal lengths and (cx, cy) is the principal point | |
| image_size_hw (tuple): Tuple of (height, width) of the image in pixels. | |
| Required for computing field of view values. For example: (256, 512). | |
| pose_encoding_type (str): Type of pose encoding to use. Currently only | |
| supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). | |
| Returns: | |
| torch.Tensor: Encoded camera pose parameters with shape BxSx9. | |
| For "absT_quaR_FoV" type, the 9 dimensions are: | |
| - [:3] = absolute translation vector T (3D) | |
| - [3:7] = rotation as quaternion quat (4D) | |
| - [7:] = field of view (2D) | |
| """ | |
| # extrinsics: BxSx3x4 | |
| # intrinsics: BxSx3x3 | |
| if pose_encoding_type == "absT_quaR_FoV": | |
| R = extrinsics[:, :, :3, :3] # BxSx3x3 | |
| T = extrinsics[:, :, :3, 3] # BxSx3 | |
| quat = mat_to_quat(R) | |
| # Note the order of h and w here | |
| H, W = image_size_hw | |
| fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1]) | |
| fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0]) | |
| pose_encoding = torch.cat([T, quat, fov_h[..., None], fov_w[..., None]], dim=-1).float() | |
| else: | |
| raise NotImplementedError | |
| return pose_encoding | |
| def pose_encoding_to_extri_intri( | |
| pose_encoding, image_size_hw=None, pose_encoding_type="absT_quaR_FoV", build_intrinsics=True # e.g., (256, 512) | |
| ): | |
| """Convert a pose encoding back to camera extrinsics and intrinsics. | |
| This function performs the inverse operation of extri_intri_to_pose_encoding, | |
| reconstructing the full camera parameters from the compact encoding. | |
| Args: | |
| pose_encoding (torch.Tensor): Encoded camera pose parameters with shape BxSx9, | |
| where B is batch size and S is sequence length. | |
| For "absT_quaR_FoV" type, the 9 dimensions are: | |
| - [:3] = absolute translation vector T (3D) | |
| - [3:7] = rotation as quaternion quat (4D) | |
| - [7:] = field of view (2D) | |
| image_size_hw (tuple): Tuple of (height, width) of the image in pixels. | |
| Required for reconstructing intrinsics from field of view values. | |
| For example: (256, 512). | |
| pose_encoding_type (str): Type of pose encoding used. Currently only | |
| supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view). | |
| build_intrinsics (bool): Whether to reconstruct the intrinsics matrix. | |
| If False, only extrinsics are returned and intrinsics will be None. | |
| Returns: | |
| tuple: (extrinsics, intrinsics) | |
| - extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4. | |
| In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world | |
| transformation. The format is [R|t] where R is a 3x3 rotation matrix and t is | |
| a 3x1 translation vector. | |
| - intrinsics (torch.Tensor or None): Camera intrinsic parameters with shape BxSx3x3, | |
| or None if build_intrinsics is False. Defined in pixels, with format: | |
| [[fx, 0, cx], | |
| [0, fy, cy], | |
| [0, 0, 1]] | |
| where fx, fy are focal lengths and (cx, cy) is the principal point, | |
| assumed to be at the center of the image (W/2, H/2). | |
| """ | |
| intrinsics = None | |
| if pose_encoding_type == "absT_quaR_FoV": | |
| T = pose_encoding[..., :3] | |
| quat = pose_encoding[..., 3:7] | |
| fov_h = pose_encoding[..., 7] | |
| fov_w = pose_encoding[..., 8] | |
| R = quat_to_mat(quat) | |
| extrinsics = torch.cat([R, T[..., None]], dim=-1) | |
| if build_intrinsics: | |
| H, W = image_size_hw | |
| fy = (H / 2.0) / torch.tan(fov_h / 2.0) | |
| fx = (W / 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] = W / 2 | |
| intrinsics[..., 1, 2] = H / 2 | |
| intrinsics[..., 2, 2] = 1.0 # Set the homogeneous coordinate to 1 | |
| elif pose_encoding_type == "absT_quaR": | |
| T = pose_encoding[..., :3] | |
| quat = pose_encoding[..., 3:7] | |
| R = quat_to_mat(quat) | |
| extrinsics = torch.cat([R, T[..., None]], dim=-1) | |
| intrinsics = None | |
| return extrinsics, intrinsics | |
| def convert_pt3d_RT_to_opencv(Rot, Trans): | |
| """ | |
| Convert Point3D extrinsic matrices to OpenCV convention. | |
| Args: | |
| Rot: 3D rotation matrix in Point3D format | |
| Trans: 3D translation vector in Point3D format | |
| Returns: | |
| extri_opencv: 3x4 extrinsic matrix in OpenCV format | |
| """ | |
| rot_pt3d = np.array(Rot) | |
| trans_pt3d = np.array(Trans) | |
| trans_pt3d[:2] *= -1 | |
| rot_pt3d[:, :2] *= -1 | |
| rot_pt3d = rot_pt3d.transpose(1, 0) | |
| extri_opencv = np.hstack((rot_pt3d, trans_pt3d[:, None])) | |
| return extri_opencv | |
| def build_pair_index(N, B=1): | |
| """ | |
| Build indices for all possible pairs of frames. | |
| Args: | |
| N: Number of frames | |
| B: Batch size | |
| Returns: | |
| i1, i2: Indices for all possible pairs | |
| """ | |
| i1_, i2_ = torch.combinations(torch.arange(N), 2, with_replacement=False).unbind(-1) | |
| i1, i2 = [(i[None] + torch.arange(B)[:, None] * N).reshape(-1) for i in [i1_, i2_]] | |
| return i1, i2 | |
| def rotation_angle(rot_gt, rot_pred, batch_size=None, eps=1e-15): | |
| """ | |
| Calculate rotation angle error between ground truth and predicted rotations. | |
| Args: | |
| rot_gt: Ground truth rotation matrices | |
| rot_pred: Predicted rotation matrices | |
| batch_size: Batch size for reshaping the result | |
| eps: Small value to avoid numerical issues | |
| Returns: | |
| Rotation angle error in degrees | |
| """ | |
| q_pred = mat_to_quat(rot_pred) | |
| q_gt = mat_to_quat(rot_gt) | |
| loss_q = (1 - (q_pred * q_gt).sum(dim=1) ** 2).clamp(min=eps) | |
| err_q = torch.arccos(1 - 2 * loss_q) | |
| rel_rangle_deg = err_q * 180 / np.pi | |
| if batch_size is not None: | |
| rel_rangle_deg = rel_rangle_deg.reshape(batch_size, -1) | |
| return rel_rangle_deg | |
| def translation_angle(tvec_gt, tvec_pred, batch_size=None, ambiguity=True): | |
| """ | |
| Calculate translation angle error between ground truth and predicted translations. | |
| Args: | |
| tvec_gt: Ground truth translation vectors | |
| tvec_pred: Predicted translation vectors | |
| batch_size: Batch size for reshaping the result | |
| ambiguity: Whether to handle direction ambiguity | |
| Returns: | |
| Translation angle error in degrees | |
| """ | |
| rel_tangle_deg = compare_translation_by_angle(tvec_gt, tvec_pred) | |
| rel_tangle_deg = rel_tangle_deg * 180.0 / np.pi | |
| if ambiguity: | |
| rel_tangle_deg = torch.min(rel_tangle_deg, (180 - rel_tangle_deg).abs()) | |
| if batch_size is not None: | |
| rel_tangle_deg = rel_tangle_deg.reshape(batch_size, -1) | |
| return rel_tangle_deg | |
| def compare_translation_by_angle(t_gt, t, eps=1e-15, default_err=1e6): | |
| """ | |
| Normalize the translation vectors and compute the angle between them. | |
| Args: | |
| t_gt: Ground truth translation vectors | |
| t: Predicted translation vectors | |
| eps: Small value to avoid division by zero | |
| default_err: Default error value for invalid cases | |
| Returns: | |
| Angular error between translation vectors in radians | |
| """ | |
| t_norm = torch.norm(t, dim=1, keepdim=True) | |
| t = t / (t_norm + eps) | |
| t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True) | |
| t_gt = t_gt / (t_gt_norm + eps) | |
| loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps) | |
| err_t = torch.acos(torch.sqrt(1 - loss_t)) | |
| err_t[torch.isnan(err_t) | torch.isinf(err_t)] = default_err | |
| return err_t | |
| def calculate_auc_np(r_error, t_error, max_threshold=30): | |
| """ | |
| Calculate the Area Under the Curve (AUC) for the given error arrays using NumPy. | |
| Args: | |
| r_error: numpy array representing R error values (Degree) | |
| t_error: numpy array representing T error values (Degree) | |
| max_threshold: Maximum threshold value for binning the histogram | |
| Returns: | |
| AUC value and the normalized histogram | |
| """ | |
| error_matrix = np.concatenate((r_error[:, None], t_error[:, None]), axis=1) | |
| max_errors = np.max(error_matrix, axis=1) | |
| bins = np.arange(max_threshold + 1) | |
| histogram, _ = np.histogram(max_errors, bins=bins) | |
| num_pairs = float(len(max_errors)) | |
| normalized_histogram = histogram.astype(float) / num_pairs | |
| return np.mean(np.cumsum(normalized_histogram)), normalized_histogram | |
| def se3_to_relative_pose_error(pred_se3, gt_se3, num_frames): | |
| """ | |
| Compute rotation and translation errors between predicted and ground truth poses. | |
| This function assumes the input poses are world-to-camera (w2c) transformations. | |
| Args: | |
| pred_se3: Predicted SE(3) transformations (w2c), shape (N, 4, 4) | |
| gt_se3: Ground truth SE(3) transformations (w2c), shape (N, 4, 4) | |
| num_frames: Number of frames (N) | |
| Returns: | |
| Rotation and translation angle errors in degrees | |
| """ | |
| pair_idx_i1, pair_idx_i2 = build_pair_index(num_frames) | |
| relative_pose_gt = gt_se3[pair_idx_i1].bmm( | |
| closed_form_inverse_se3(gt_se3[pair_idx_i2]) | |
| ) | |
| relative_pose_pred = pred_se3[pair_idx_i1].bmm( | |
| closed_form_inverse_se3(pred_se3[pair_idx_i2]) | |
| ) | |
| rel_rangle_deg = rotation_angle( | |
| relative_pose_gt[:, :3, :3], relative_pose_pred[:, :3, :3] | |
| ) | |
| rel_tangle_deg = translation_angle( | |
| relative_pose_gt[:, :3, 3], relative_pose_pred[:, :3, 3] | |
| ) | |
| return rel_rangle_deg, rel_tangle_deg | |
| def colmap_to_opencv_intrinsics(K): | |
| """ | |
| Modify camera intrinsics to follow a different convention. | |
| Coordinates of the center of the top-left pixels are by default: | |
| - (0.5, 0.5) in Colmap | |
| - (0,0) in OpenCV | |
| """ | |
| K = K.copy() | |
| K[..., 0, 2] -= 0.5 | |
| K[..., 1, 2] -= 0.5 | |
| return K | |
| def read_camera_parameters(filename): | |
| with open(filename) as f: | |
| lines = f.readlines() | |
| lines = [line.rstrip() for line in lines] | |
| # extrinsics: line [1,5), 4x4 matrix | |
| extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ').reshape((4, 4)) | |
| # intrinsics: line [7-10), 3x3 matrix | |
| intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ').reshape((3, 3)) | |
| return intrinsics, extrinsics |