Source code for ct.convert

"""
Functions for converting between different camera parameters and representations.
"""

import cv2
import numpy as np
import open3d as o3d
from jaxtyping import Float
from typing import Optional, Tuple

from . import sanity
from . import convert


[docs] def pad_0001(array): """ Pad [0, 0, 0, 1] to the bottom row. Args: array: (3, 4) or (N, 3, 4). Returns: Array of shape (4, 4) or (N, 4, 4). """ if array.ndim == 2: if not array.shape == (3, 4): raise ValueError(f"Expected array of shape (3, 4), but got {array.shape}.") elif array.ndim == 3: if not array.shape[-2:] == (3, 4): raise ValueError( f"Expected array of shape (N, 3, 4), but got {array.shape}." ) else: raise ValueError( f"Expected array of shape (3, 4) or (N, 3, 4), but got {array.shape}." ) if array.ndim == 2: bottom = np.array([0, 0, 0, 1], dtype=array.dtype) return np.concatenate([array, bottom[None, :]], axis=0) elif array.ndim == 3: bottom_single = np.array([0, 0, 0, 1], dtype=array.dtype) bottom = np.broadcast_to(bottom_single, (array.shape[0], 1, 4)) return np.concatenate([array, bottom], axis=-2) else: raise ValueError("Should not reach here.")
[docs] def rm_pad_0001(array, check_vals=False): """ Remove the bottom row of [0, 0, 0, 1]. Args: array: (4, 4) or (N, 4, 4). check_vals (bool): If True, check that the bottom row is [0, 0, 0, 1]. Returns: Array of shape (3, 4) or (N, 3, 4). """ # Check shapes. if array.ndim == 2: if not array.shape == (4, 4): raise ValueError(f"Expected array of shape (4, 4), but got {array.shape}.") elif array.ndim == 3: if not array.shape[-2:] == (4, 4): raise ValueError( f"Expected array of shape (N, 4, 4), but got {array.shape}." ) else: raise ValueError( f"Expected array of shape (4, 4) or (N, 4, 4), but got {array.shape}." ) # Check vals. if check_vals: if array.ndim == 2: bottom = array[3, :] if not np.allclose(bottom, [0, 0, 0, 1]): raise ValueError( f"Expected bottom row to be [0, 0, 0, 1], but got {bottom}." ) elif array.ndim == 3: bottom = array[:, 3:4, :] expected_bottom = np.broadcast_to([0, 0, 0, 1], (array.shape[0], 1, 4)) if not np.allclose(bottom, expected_bottom): raise ValueError( f"Expected bottom row to be {expected_bottom}, but got {bottom}." ) else: raise ValueError("Should not reach here.") return array[..., :3, :]
[docs] def to_homo(array): """ Convert a 2D array to homogeneous coordinates by appending a column of ones. Args: array: A 2D numpy array of shape (N, M). Returns: A numpy array of shape (N, M+1) with a column of ones appended. """ if not isinstance(array, np.ndarray) or array.ndim != 2: raise ValueError(f"Input must be a 2D numpy array, but got {array.shape}.") ones = np.ones((array.shape[0], 1), dtype=array.dtype) return np.hstack((array, ones))
[docs] def from_homo(array): """ Convert an array from homogeneous to Cartesian coordinates by dividing by the last column and removing it. Args: array: A 2D numpy array of shape (N, M) in homogeneous coordinates. Returns: A numpy array of shape (N, M-1) in Cartesian coordinates. """ if not isinstance(array, np.ndarray) or array.ndim != 2: raise ValueError(f"Input must be a 2D numpy array, but got {array.shape}.") if array.shape[1] < 2: raise ValueError( f"Input array must have at least two columns for removing " f"homogeneous coordinate, but got shape {array.shape}." ) return array[:, :-1] / array[:, -1, np.newaxis]
[docs] def R_to_quat(R): """ Convert rotation matrix to quaternion. Args: R: Rotation matrix of shape (3, 3) or (N, 3, 3). Returns: Quaternion of shape (4,) or (N, 4). The quaternion is normalized. """ # https://github.com/isl-org/StableViewSynthesis/tree/main/co R = R.reshape(-1, 3, 3) q = np.empty((R.shape[0], 4), dtype=R.dtype) q[:, 0] = np.sqrt(np.maximum(0, 1 + R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2])) q[:, 1] = np.sqrt(np.maximum(0, 1 + R[:, 0, 0] - R[:, 1, 1] - R[:, 2, 2])) q[:, 2] = np.sqrt(np.maximum(0, 1 - R[:, 0, 0] + R[:, 1, 1] - R[:, 2, 2])) q[:, 3] = np.sqrt(np.maximum(0, 1 - R[:, 0, 0] - R[:, 1, 1] + R[:, 2, 2])) q[:, 1] *= 2 * (R[:, 2, 1] > R[:, 1, 2]) - 1 q[:, 2] *= 2 * (R[:, 0, 2] > R[:, 2, 0]) - 1 q[:, 3] *= 2 * (R[:, 1, 0] > R[:, 0, 1]) - 1 q /= np.linalg.norm(q, axis=1, keepdims=True) return q.squeeze()
[docs] def T_to_C(T: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "3"]: """ Convert extrinsic matrix T to camera center C. Args: T: Extrinsic matrix (world-to-camera) of shape (4, 4). Returns: Camera center in world coordinates of shape (3,). """ sanity.assert_T(T) R, t = T[:3, :3], T[:3, 3] return R_t_to_C(R, t)
[docs] def pose_to_C(pose: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "3"]: """ Convert pose matrix to camera center C. Args: pose: Pose matrix (camera-to-world) of shape (4, 4). Returns: Camera center in world coordinates of shape (3,). """ sanity.assert_pose(pose) C = pose[:3, 3] return C
[docs] def T_to_pose(T): """ Convert extrinsic matrix T to pose matrix. Args: T: Extrinsic matrix (world-to-camera) of shape (4, 4). Returns: Pose matrix (camera-to-world) of shape (4, 4), which is the inverse of T. """ sanity.assert_T(T) return np.linalg.inv(T)
[docs] def pose_to_T(pose): """ Convert pose matrix to extrinsic matrix T. Args: pose: Pose matrix (camera-to-world) of shape (4, 4). Returns: Extrinsic matrix (world-to-camera) of shape (4, 4), which is the inverse of pose. """ sanity.assert_T(pose) return np.linalg.inv(pose)
[docs] def T_opengl_to_opencv(T: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "4 4"]: """ Convert T from OpenGL convention to OpenCV convention. - OpenCV - +X: Right - +Y: Down - +Z: The view direction, pointing forward and away from the camera - Used in: OpenCV, COLMAP, camtools default - OpenGL - +X: Right - +Y: Up - +Z: The negative view direction, pointing back and away from the camera - -Z: The view direction - Used in: OpenGL, Blender, Nerfstudio https://docs.nerf.studio/quickstart/data_conventions.html#coordinate-conventions Args: T: Extrinsic matrix (world-to-camera) of shape (4, 4) in OpenCV convention. Returns: Extrinsic matrix (world-to-camera) of shape (4, 4) in OpenGL convention. """ sanity.assert_T(T) # pose = T_to_pose(T) # pose = pose_opengl_to_opencv(pose) # T = pose_to_T(pose) T = np.copy(T) T[1:3, 0:4] *= -1 T = T[:, [1, 0, 2, 3]] T[:, 2] *= -1 return T
[docs] def T_opencv_to_opengl(T: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "4 4"]: """ Convert T from OpenCV convention to OpenGL convention. - OpenCV - +X: Right - +Y: Down - +Z: The view direction, pointing forward and away from the camera - Used in: OpenCV, COLMAP, camtools default - OpenGL - +X: Right - +Y: Up - +Z: The negative view direction, pointing back and away from the camera - -Z: The view direction - Used in: OpenGL, Blender, Nerfstudio https://docs.nerf.studio/quickstart/data_conventions.html#coordinate-conventions Args: T: Extrinsic matrix (world-to-camera) of shape (4, 4) in OpenCV convention. Returns: Extrinsic matrix (world-to-camera) of shape (4, 4) in OpenGL convention. """ sanity.assert_T(T) # pose = T_to_pose(T) # pose = pose_opencv_to_opengl(pose) # T = pose_to_T(pose) T = np.copy(T) T[:, 2] *= -1 T = T[:, [1, 0, 2, 3]] T[1:3, 0:4] *= -1 return T
[docs] def pose_opengl_to_opencv(pose: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "4 4"]: """ Convert pose from OpenGL convention to OpenCV convention. - OpenCV - +X: Right - +Y: Down - +Z: The view direction, pointing forward and away from the camera - Used in: OpenCV, COLMAP, camtools default - OpenGL - +X: Right - +Y: Up - +Z: The negative view direction, pointing back and away from the camera - -Z: The view direction - Used in: OpenGL, Blender, Nerfstudio https://docs.nerf.studio/quickstart/data_conventions.html#coordinate-conventions Args: pose: Pose matrix (camera-to-world) of shape (4, 4) in OpenGL convention. Returns: Pose matrix (camera-to-world) of shape (4, 4) in OpenCV convention. """ sanity.assert_pose(pose) pose = np.copy(pose) pose[2, :] *= -1 pose = pose[[1, 0, 2, 3], :] pose[0:3, 1:3] *= -1 return pose
[docs] def pose_opencv_to_opengl(pose: Float[np.ndarray, "4 4"]) -> Float[np.ndarray, "4 4"]: """ Convert pose from OpenCV convention to OpenGL convention. - OpenCV - +X: Right - +Y: Down - +Z: The view direction, pointing forward and away from the camera - Used in: OpenCV, COLMAP, camtools default - OpenGL - +X: Right - +Y: Up - +Z: The negative view direction, pointing back and away from the camera - -Z: The view direction - Used in: OpenGL, Blender, Nerfstudio https://docs.nerf.studio/quickstart/data_conventions.html#coordinate-conventions Args: pose: Pose matrix (camera-to-world) of shape (4, 4) in OpenCV convention. Returns: Pose matrix (camera-to-world) of shape (4, 4) in OpenGL convention. """ sanity.assert_pose(pose) pose = np.copy(pose) pose[0:3, 1:3] *= -1 pose = pose[[1, 0, 2, 3], :] pose[2, :] *= -1 return pose
[docs] def R_t_to_C( R: Float[np.ndarray, "3 3"], t: Float[np.ndarray, "3"], ) -> Float[np.ndarray, "3"]: """ Convert rotation matrix R and translation vector t to camera center C. Args: R: Rotation matrix of shape (3, 3). t: Translation vector of shape (3,). Returns: Camera center in world coordinates of shape (3,). """ # Equivalently, # C = - R.T @ t # C = - np.linalg.inv(R) @ t # C = pose[:3, 3] = np.linalg.inv(R_t_to_T(R, t))[:3, 3] t = t.reshape(-1, 3, 1) R = R.reshape(-1, 3, 3) C = -R.transpose(0, 2, 1) @ t return C.squeeze()
[docs] def R_C_to_t( R: Float[np.ndarray, "3 3"], C: Float[np.ndarray, "3"], ) -> Float[np.ndarray, "3"]: """ Convert rotation matrix R and camera center C to translation vector t. Args: R: Rotation matrix of shape (3, 3) or (N, 3, 3). C: Camera center in world coordinates of shape (3,) or (N, 3). Returns: Translation vector of shape (3,) or (N, 3). """ # https://github.com/isl-org/StableViewSynthesis/blob/main/data/create_custom_track.py C = C.reshape(-1, 3, 1) R = R.reshape(-1, 3, 3) t = -R @ C return t.squeeze()
[docs] def roll_pitch_yaw_to_R( roll: float, pitch: float, yaw: float, ) -> Float[np.ndarray, "3 3"]: """ Convert roll, pitch, and yaw angles to a rotation matrix. Args: roll: Rotation around the x-axis in radians. pitch: Rotation around the y-axis in radians. yaw: Rotation around the z-axis in radians. Returns: Rotation matrix of shape (3, 3). """ rx_roll = np.array( [ [1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)], ] ) ry_pitch = np.array( [ [np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)], ] ) rz_yaw = np.array( [ [np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1], ] ) R = rz_yaw @ ry_pitch @ rx_roll return R
[docs] def R_t_to_T( R: Float[np.ndarray, "3 3"], t: Float[np.ndarray, "3"], ) -> Float[np.ndarray, "4 4"]: """ Convert rotation matrix R and translation vector t to extrinsic matrix T. Args: R: Rotation matrix of shape (3, 3). t: Translation vector of shape (3,). Returns: Extrinsic matrix (world-to-camera) of shape (4, 4). """ T = np.eye(4) T[:3, :3] = R T[:3, 3] = t return T
[docs] def T_to_R_t( T: Float[np.ndarray, "4 4"], ) -> Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "3"]]: """ Decompose extrinsic matrix T into rotation matrix R and translation vector t. Args: T: Extrinsic matrix (world-to-camera) of shape (4, 4). Returns: Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "3"]]: - R: Rotation matrix of shape (3, 3) - t: Translation vector of shape (3,) """ sanity.assert_T(T) R = T[:3, :3] t = T[:3, 3] return R, t
[docs] def P_to_K_R_t( P: Float[np.ndarray, "3 4"], ) -> Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "3 3"], Float[np.ndarray, "3"]]: """ Decompose projection matrix P into intrinsic matrix K, rotation matrix R, and translation vector t. Args: P: Projection matrix of shape (3, 4). Returns: Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "3 3"], Float[np.ndarray, "3"]]: - K: Intrinsic matrix of shape (3, 3) - R: Rotation matrix of shape (3, 3) - t: Translation vector of shape (3,) """ ( camera_matrix, rot_matrix, trans_vect, rot_matrix_x, rot_matrix_y, rot_matrix_z, euler_angles, ) = cv2.decomposeProjectionMatrix(P) K = camera_matrix K = K / K[2, 2] R = rot_matrix t = -rot_matrix @ (trans_vect[:3] / trans_vect[3]) return K, R, t.squeeze()
[docs] def P_to_K_T( P: Float[np.ndarray, "3 4"], ) -> Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "4 4"]]: """ Decompose projection matrix P into intrinsic matrix K and extrinsic matrix T. Args: P: Projection matrix of shape (3, 4). Returns: Tuple[Float[np.ndarray, "3 3"], Float[np.ndarray, "4 4"]]: - K: Intrinsic matrix of shape (3, 3) - T: Extrinsic matrix (world-to-camera) of shape (4, 4) """ K, R, t = P_to_K_R_t(P) T = R_t_to_T(R, t) return K, T
[docs] def K_T_to_P( K: Float[np.ndarray, "3 3"], T: Float[np.ndarray, "4 4"], ) -> Float[np.ndarray, "3 4"]: """ Compute projection matrix P from intrinsic matrix K and extrinsic matrix T. Args: K: Intrinsic matrix of shape (3, 3). T: Extrinsic matrix (world-to-camera) of shape (4, 4). Returns: Projection matrix of shape (3, 4). """ return K @ T[:3, :]
[docs] def K_R_t_to_P( K: Float[np.ndarray, "3 3"], R: Float[np.ndarray, "3 3"], t: Float[np.ndarray, "3"], ) -> Float[np.ndarray, "3 4"]: """ Compute projection matrix from intrinsic matrix, rotation matrix, and translation vector. Args: K: Intrinsic matrix of shape (3, 3). R: Rotation matrix of shape (3, 3). t: Translation vector of shape (3,). Returns: Projection matrix of shape (3, 4). """ T = R_t_to_T(R, t) P = K @ T[:3, :] return P
[docs] def K_R_t_to_W2P( K: Float[np.ndarray, "3 3"], R: Float[np.ndarray, "3 3"], t: Float[np.ndarray, "3"], ) -> Float[np.ndarray, "4 4"]: """ Compute world-to-projection matrix from intrinsic matrix, rotation matrix, and translation vector. Args: K: Intrinsic matrix of shape (3, 3). R: Rotation matrix of shape (3, 3). t: Translation vector of shape (3,). Returns: World-to-projection matrix of shape (4, 4). """ return P_to_W2P(K_R_t_to_P(K, R, t))
[docs] def K_T_to_W2P( K: Float[np.ndarray, "3 3"], T: Float[np.ndarray, "4 4"], ) -> Float[np.ndarray, "4 4"]: """ Compute world-to-projection matrix from intrinsic matrix and transformation matrix. Args: K: Intrinsic matrix of shape (3, 3). T: Extrinsic matrix (world-to-camera matrix) of shape (4, 4). Returns: World-to-projection matrix of shape (4, 4). """ R, t = T_to_R_t(T) return K_R_t_to_W2P(K, R, t)
[docs] def P_to_W2P( P: Float[np.ndarray, "3 4"], ) -> Float[np.ndarray, "4 4"]: """ Convert projection matrix to world-to-projection matrix. Args: P: Projection matrix of shape (3, 4). Returns: World-to-projection matrix of shape (4, 4). """ sanity.assert_shape_3x4(P, name="P") W2P = convert.pad_0001(P) return W2P
[docs] def W2P_to_P( W2P: Float[np.ndarray, "4 4"], ) -> Float[np.ndarray, "3 4"]: """ Convert world-to-projection matrix to projection matrix. Args: W2P: World-to-projection matrix of shape (4, 4). Returns: Projection matrix of shape (3, 4). """ if W2P.shape != (4, 4): raise ValueError(f"Expected W2P of shape (4, 4), but got {W2P.shape}.") P = convert.rm_pad_0001(W2P, check_vals=True) return P
[docs] def fx_fy_cx_cy_to_K( fx: float, fy: float, cx: float, cy: float, ) -> Float[np.ndarray, "3 3"]: """ Create intrinsic matrix from focal lengths and principal point coordinates. Args: fx: Focal length in x direction. fy: Focal length in y direction. cx: Principal point x coordinate. cy: Principal point y coordinate. Returns: Intrinsic matrix of shape (3, 3). """ K = np.zeros((3, 3)) K[0, 0] = fx K[1, 1] = fy K[0, 2] = cx K[1, 2] = cy K[2, 2] = 1 return K
[docs] def K_to_fx_fy_cx_cy( K: Float[np.ndarray, "3 3"], ) -> Tuple[float, float, float, float]: """ Extract focal lengths and principal point coordinates from intrinsic matrix. Args: K: Intrinsic matrix of shape (3, 3). Returns: Tuple[float, float, float, float]: - fx: Focal length in x direction - fy: Focal length in y direction - cx: Principal point x coordinate - cy: Principal point y coordinate """ fx = K[0, 0] fy = K[1, 1] cx = K[0, 2] cy = K[1, 2] # <class 'numpy.float64'> to <class 'float'> return float(fx), float(fy), float(cx), float(cy)
[docs] def euler_to_R( yaw: float, pitch: float, roll: float, ) -> Float[np.ndarray, "3 3"]: """ Convert Euler angles to rotation matrix. Given a unit vector x, R @ x is x rotated by applying yaw, pitch, and roll consecutively. Ref: https://en.wikipedia.org/wiki/Euler_angles Args: yaw (float): Rotation around the z-axis (from x-axis to y-axis). pitch (float): Rotation around the y-axis (from z-axis to x-axis). roll (float): Rotation around the x-axis (from y-axis to z-axis). Returns: Rotation matrix R of shape (3, 3). """ sin_y = np.sin(yaw) cos_y = np.cos(yaw) sin_p = np.sin(pitch) cos_p = np.cos(pitch) sin_r = np.sin(roll) cos_r = np.cos(roll) R = np.array( [ [ cos_y * cos_p, cos_y * sin_p * sin_r - sin_y * cos_r, cos_y * sin_p * cos_r + sin_y * sin_r, ], [ sin_y * cos_p, sin_y * sin_p * sin_r + cos_y * cos_r, sin_y * sin_p * cos_r - cos_y * sin_r, ], [ -sin_p, cos_p * sin_r, cos_p * cos_r, ], ] ) return R
[docs] def spherical_to_T_towards_origin( radius: float, theta: float, phi: float, ) -> Float[np.ndarray, "4 4"]: """ Convert spherical coordinates (ISO convention) to T, where the cameras looks at the origin from a distance (radius), and the camera up direction alines with the z-axis (the angle between the up direction and the z-axis is smaller than pi/2). Args: radius (float): Distance from the origin. theta (float): Inclination, angle w.r.t. positive polar (+z) axis. Range: [0, pi]. phi (float): Azimuth, rotation angle from the initial meridian (x-y) plane. Range: [0, 2*pi]. Returns: T of shape (4, 4). Ref: https://en.wikipedia.org/wiki/Spherical_coordinate_system """ if not 0 <= theta <= np.pi: raise ValueError(f"Expected theta in [0, pi], but got {theta}.") if not 0 <= phi <= 2 * np.pi: raise ValueError(f"Expected phi in [0, 2*pi], but got {phi}.") x = radius * np.sin(theta) * np.cos(phi) y = radius * np.sin(theta) * np.sin(phi) z = radius * np.cos(theta) # Default : look at +Z, up is -Y. # After init: look at +X, up is +Z. init_R = euler_to_R(-np.pi / 2, 0, -np.pi / 2) # Rotate along z axis. phi_R = euler_to_R(phi + np.pi, 0, 0) # Rotate along y axis. theta_R = euler_to_R(0, np.pi / 2 - theta, 0) # Combine rotations, the order matters. R = phi_R @ theta_R @ init_R pose = np.eye(4) pose[:3, :3] = R pose[:3, 3] = [x, y, z] T = pose_to_T(pose) return T
[docs] def mesh_to_lineset( mesh: o3d.geometry.TriangleMesh, color: Optional[Float[np.ndarray, "3"]] = None, ) -> o3d.geometry.LineSet: """ Convert Open3D mesh to Open3D lineset. """ if not isinstance(mesh, o3d.geometry.TriangleMesh): raise ValueError(f"Expected Open3D mesh, but got {type(mesh)}.") vertices = np.asarray(mesh.vertices) triangles = np.asarray(mesh.triangles) edges = set() for triangle in triangles: edges.add(tuple(sorted([triangle[0], triangle[1]]))) edges.add(tuple(sorted([triangle[1], triangle[2]]))) edges.add(tuple(sorted([triangle[2], triangle[0]]))) edges = np.array(list(edges)) lineset = o3d.geometry.LineSet() lineset.points = o3d.utility.Vector3dVector(vertices) lineset.lines = o3d.utility.Vector2iVector(edges) if color is not None: if len(color) != 3: raise ValueError(f"Expected color of shape (3,), but got {color.shape}.") lineset.paint_uniform_color(color) return lineset
[docs] def im_distance_to_im_depth( im_distance: Float[np.ndarray, "h w"], K: Float[np.ndarray, "3 3"], ) -> Float[np.ndarray, "h w"]: """ Convert distance image to depth image. Args: im_distance: Distance image (H, W), float. K: Camera intrinsic matrix (3, 3). Returns: Depth image (H, W), float. """ if not im_distance.ndim == 2: raise ValueError( f"Expected im_distance of shape (H, W), but got {im_distance.shape}." ) sanity.assert_K(K) height, width = im_distance.shape fx, fy = K[0, 0], K[1, 1] cx, cy = K[0, 2], K[1, 2] dtype = im_distance.dtype u = np.arange(width) v = np.arange(height) u_grid, v_grid = np.meshgrid(u, v) u_norm = (u_grid - cx) / fx v_norm = (v_grid - cy) / fy norm_square = u_norm**2 + v_norm**2 im_depth = im_distance / np.sqrt(norm_square + 1) im_depth = im_depth.astype(dtype) return im_depth
[docs] def im_depth_to_im_distance( im_depth: Float[np.ndarray, "h w"], K: Float[np.ndarray, "3 3"], ) -> Float[np.ndarray, "h w"]: """ Convert depth image to distance image. Args: im_depth: Depth image (H, W), float. K: Camera intrinsic matrix (3, 3). Returns: Distance image (H, W), float. """ if not im_depth.ndim == 2: raise ValueError( f"Expected im_depth of shape (H, W), but got {im_depth.shape}." ) sanity.assert_K(K) height, width = im_depth.shape fx, fy = K[0, 0], K[1, 1] cx, cy = K[0, 2], K[1, 2] dtype = im_depth.dtype u = np.arange(width) v = np.arange(height) u_grid, v_grid = np.meshgrid(u, v) u_norm = (u_grid - cx) / fx v_norm = (v_grid - cy) / fy norm_square = u_norm**2 + v_norm**2 im_distance = im_depth * np.sqrt(norm_square + 1) im_distance = im_distance.astype(dtype) return im_distance