From 67563590fc8522b2e838fc6ad1949f5790f4c918 Mon Sep 17 00:00:00 2001 From: Yixing Lao Date: Tue, 9 Jul 2024 22:19:49 +0800 Subject: [PATCH] feat: apply backend wrappers to camtools functions --- camtools/camera.py | 130 +++++++++----- camtools/colormap.py | 64 +++---- camtools/convert.py | 163 ++++++++++-------- camtools/metric.py | 9 +- camtools/sanity.py | 6 - test/test_convert.py | 394 +++++++++++++++++++++++++++++++++++++------ 6 files changed, 556 insertions(+), 210 deletions(-) diff --git a/camtools/camera.py b/camtools/camera.py index e74430f..3173af0 100644 --- a/camtools/camera.py +++ b/camtools/camera.py @@ -1,30 +1,36 @@ import open3d as o3d import numpy as np + from . import convert from . import sanity from . import solver +from .backend import Tensor, tensor_backend_numpy +from jaxtyping import Float +from typing import List, Tuple, Dict + +@tensor_backend_numpy def create_camera_frustums( - Ks, - Ts, - image_whs=None, - size=0.1, - color=(0, 0, 1), - highlight_color_map=None, - center_line=True, - center_line_color=(1, 0, 0), - up_triangle=True, - center_ray=False, + Ks: Float[Tensor, "n 3 3"], + Ts: Float[Tensor, "n 4 4"], + image_whs: List[Tuple[int, int]] = None, + size: float = 0.1, + color: Float[Tensor, "3"] = (0.0, 0.0, 1.0), + highlight_color_map: Dict[int, Float[Tensor, "3"]] = None, + center_line: bool = True, + center_line_color: Float[Tensor, "3"] = (1.0, 0.0, 0.0), + up_triangle: bool = True, + center_ray: bool = False, ): """ Create camera frustums in lineset. Args: - Ks: List of 3x3 camera intrinsics matrices. You can set Ks to None if - the intrinsics are not available. In this case, a dummy intrinsics + Ks: List of 3x3 camera intrinsics. You can set Ks to None if the + intrinsics are not available. In this case, a dummy intrinsics matrix will be used. - Ts: List of 4x4 camera extrinsics matrices. + Ts: List of 4x4 camera camera extrinsics. image_whs: List of image width and height. If None, the image width and height are determined from the camera intrinsics by assuming that the camera offset is exactly at the center of the image. @@ -33,8 +39,8 @@ def create_camera_frustums( highlight_color_map: A map of camera_index to color, specifying the colors of the highlighted cameras. Index wrapping is supported. For example, to highlight the start and stop cameras, use: - highlight_color_map = {0: [0, 1, 0], -1: [1, 0, 0]}. If None, no - camera is highlighted. + highlight_color_map = {0: [0.0, 1.0, 0.0], -1: [1.0, 0.0, 0.0]}. + If None, no camera is highlighted. center_line: If True, the camera center line will be drawn. center_line_color: Color of the camera center line. up_triangle: If True, the up triangle will be drawn. @@ -116,18 +122,37 @@ def create_camera_frustums( def create_camera_frustum_with_Ts( - Ts, - image_whs=None, - size=0.1, - color=(0, 0, 1), - highlight_color_map=None, - center_line=True, - center_line_color=(1, 0, 0), - up_triangle=True, - center_ray=False, + Ts: Float[Tensor, "n 4 4"], + image_whs: List[Tuple[int, int]] = None, + size: float = 0.1, + color: Float[Tensor, "3"] = (0.0, 0.0, 1.0), + highlight_color_map: Dict[int, Float[Tensor, "3"]] = None, + center_line: bool = True, + center_line_color: Float[Tensor, "3"] = (1.0, 0.0, 0.0), + up_triangle: bool = True, + center_ray: bool = False, ): """ - Returns ct.camera.create_camera_frustums(Ks=None, Ts, ...). + Create camera frustums in lineset. Returns + ct.camera.create_camera_frustums(Ks=None, Ts, ...). + + Args: + Ts: List of 4x4 camera camera extrinsics. + image_whs: List of image width and height. If None, the image width and + height are determined from the camera intrinsics by assuming that + the camera offset is exactly at the center of the image. + size: Distance from the camera center to image plane in world coordinates. + color: Color of the camera frustums. + highlight_color_map: A map of camera_index to color, specifying the + colors of the highlighted cameras. Index wrapping is supported. + For example, to highlight the start and stop cameras, use: + highlight_color_map = {0: [0.0, 1.0, 0.0], -1: [1.0, 0.0, 0.0]}. + If None, no camera is highlighted. + center_line: If True, the camera center line will be drawn. + center_line_color: Color of the camera center line. + up_triangle: If True, the up triangle will be drawn. + center_ray: If True, the ray from camera center to the center pixel in + the image plane will be drawn. """ return create_camera_frustums( Ks=None, @@ -143,7 +168,21 @@ def create_camera_frustum_with_Ts( ) -def create_camera_center_line(Ts, color=np.array([1, 0, 0])): +@tensor_backend_numpy +def create_camera_center_line( + Ts: Float[Tensor, "n 4 4"], + color: Float[Tensor, "3"] = (1.0, 0.0, 0.0), +): + """ + Create camera center lines in lineset. + + Args: + Ts: List of 4x4 camera camera extrinsics. + color: Color of the camera center lines. + + Return: + An Open3D lineset containing all the camera center lines. + """ num_nodes = len(Ts) camera_centers = [convert.T_to_C(T) for T in Ts] @@ -158,22 +197,14 @@ def create_camera_center_line(Ts, color=np.array([1, 0, 0])): def _create_camera_frustum( - K, - T, - image_wh, - size, - color, - up_triangle, - center_ray, + K: Float[Tensor, "3 3"], + T: Float[Tensor, "4 4"], + image_wh: Tuple[int, int], + size: float, + color: Float[Tensor, "3"], + up_triangle: bool, + center_ray: bool, ): - """ - K: (3, 3) - T: (4, 4) - image:_wh: (2,) - size: float - up_triangle: bool - center_ray: bool - """ T, K, color = np.asarray(T), np.asarray(K), np.asarray(color) sanity.assert_T(T) sanity.assert_K(K) @@ -280,7 +311,22 @@ def points_2d_to_3d_world(points_2d): return ls -def _wrap_dim(dim: int, max_dim: int, inclusive: bool = False) -> int: +def _wrap_dim( + dim: int, + max_dim: int, + inclusive: bool = False, +) -> int: + """ + Wrap the dimension index to the range [0, max_dim) or [0, max_dim]. + + Args: + dim: The input dimension index. + max_dim: The maximum dimension index. + inclusive: If True, the maximum dimension index is inclusive. + + Return: + The wrapped dimension index. + """ if max_dim <= 0: raise ValueError(f"max_dim {max_dim} must be > 0.") min = -max_dim diff --git a/camtools/colormap.py b/camtools/colormap.py index 6c35cab..169c9e8 100644 --- a/camtools/colormap.py +++ b/camtools/colormap.py @@ -1,35 +1,47 @@ -import matplotlib import numpy as np from . import io +from .backend import Tensor, tensor_backend_numpy, tensor_backend_auto, ivy +from jaxtyping import Float +from typing import List, Tuple, Dict +from matplotlib import pyplot as plt -def query(points, colormap="viridis"): + +@tensor_backend_numpy +def query( + values: Float[Tensor, "..."], + colormap="viridis", +) -> Float[Tensor, "... 3"]: """ Query matplotlib's color map. Args: - points: Numpy array in float32 or float64. Valid range is [0, 1]. + values: Numpy array in float32 or float64. Valid range is [0, 1]. It + can be of arbitrary shape. As matplotlib color maps have different + behaviors for float and int, the input array should be in float. colormap: Name of matplotlib color map. Returns: - Numpy array of shape (**points.shape, 3) with dtype float32. + Numpy array of shape (**values.shape, 3) with dtype float32. """ - assert isinstance(points, np.ndarray) - - if not points.dtype == np.float32 and not points.dtype == np.float64: - raise ValueError( - "Matplotlib's colormap has different behavior for ints and floats. " - "To unify behavior, we require floats (between 0-1 if valid). " - f"However, dtype of {points.dtype} is used." - ) + try: + cmap = plt.get_cmap(colormap) + except AttributeError: + cmap = plt.colormaps.get_cmap(colormap) - cmap = matplotlib.cm.get_cmap(colormap) - colors = cmap(points)[..., :3] # Remove alpha. + colors = cmap(values) + colors = colors[..., :3] # Remove alpha channel if present return colors.astype(np.float32) -def normalize(array, vmin=0.0, vmax=1.0, clip=False): +@tensor_backend_auto +def normalize( + array: Float[Tensor, "..."], + vmin: float = 0.0, + vmax: float = 1.0, + clip: bool = False, +): """ Normalize array to [vmin, vmax]. @@ -43,30 +55,10 @@ def normalize(array, vmin=0.0, vmax=1.0, clip=False): Normalized array of the same shape as the input array. """ if clip: - array = np.clip(array, vmin, vmax) + array = ivy.clip(array, x_min=vmin, x_max=vmax) else: amin = array.min() amax = array.max() array = (array - amin) / (amax - amin) * (vmax - vmin) + vmin return array - - -def main(): - """ - Test create color map image. - """ - height = 200 - width = 1600 - - colors = query(np.linspace(0, 1, num=width)) - im = np.zeros((height, width, 3), dtype=np.float32) - for i in range(width): - im[:, i : i + 1, :] = colors[i] - - im_path = "colormap.png" - io.imwrite(im_path, im) - - -if __name__ == "__main__": - main() diff --git a/camtools/convert.py b/camtools/convert.py index 32a7fe8..275c2e6 100644 --- a/camtools/convert.py +++ b/camtools/convert.py @@ -4,42 +4,55 @@ from . import sanity from . import convert - -def pad_0001(array): +from .backend import ( + Tensor, + tensor_backend_numpy, + tensor_backend_auto, + is_torch_available, + ivy, + torch, + create_array, + create_ones, + create_empty, + get_tensor_backend, +) +from jaxtyping import Float +from typing import List, Tuple, Dict, Union +from matplotlib import pyplot as plt + + +@tensor_backend_auto +def pad_0001( + array: Union[Float[Tensor, "3 4"], Float[Tensor, "N 3 4"]] +) -> Union[Float[Tensor, "4 4"], Float[Tensor, "N 4 4"]]: """ Pad [0, 0, 0, 1] to the bottom row. Args: - array: (3, 4) or (N, 3, 4). + array: NumPy or Torch array of shape (3, 4) or (N, 3, 4). Returns: - Array of shape (4, 4) or (N, 4, 4). + NumPy or Torch 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}." - ) + dtype = array.dtype + backend = get_tensor_backend(array) + bottom = create_array([0, 0, 0, 1], dtype=dtype, backend=backend) if array.ndim == 2: - bottom = np.array([0, 0, 0, 1], dtype=array.dtype) - return np.concatenate([array, bottom[None, :]], axis=0) + return ivy.concat([array, ivy.expand_dims(bottom, axis=0)], 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) + bottom = ivy.expand_dims(bottom, axis=0) + bottom = ivy.broadcast_to(bottom, (array.shape[0], 1, 4)) + return ivy.concat([array, bottom], axis=1) else: - raise ValueError("Should not reach here.") + raise ValueError("Input array must be 2D or 3D.") -def rm_pad_0001(array, check_vals=False): +@tensor_backend_auto +def rm_pad_0001( + array: Union[Float[Tensor, "4 4"], Float[Tensor, "N 4 4"]], + check_vals: bool = False, +) -> Union[Float[Tensor, "3 4"], Float[Tensor, "N 3 4"]]: """ Remove the bottom row of [0, 0, 0, 1]. @@ -50,42 +63,29 @@ def rm_pad_0001(array, check_vals=False): 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: + backend = get_tensor_backend(array) + dtype = array.dtype + gt_bottom = create_array([0, 0, 0, 1], dtype=dtype, backend=backend) + 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.") + raise ValueError(f"Invalid array shape {array.shape}.") + + if not ivy.allclose(bottom, gt_bottom): + raise ValueError( + f"Expected bottom row to be {gt_bottom}, but got {bottom}." + ) return array[..., :3, :] -def to_homo(array): +@tensor_backend_auto +def to_homo(array: Float[Tensor, "n m"]) -> Float[Tensor, "n m+1"]: """ Convert a 2D array to homogeneous coordinates by appending a column of ones. @@ -95,48 +95,68 @@ def to_homo(array): 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)) + backend = get_tensor_backend(array) + ones = create_ones((array.shape[0], 1), dtype=array.dtype, backend=backend) + return ivy.concat([array, ones], axis=1) -def from_homo(array): +@tensor_backend_auto +def from_homo(array: Float[Tensor, "n m"]) -> Float[Tensor, "n m-1"]: """ 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. + array: A 2D array of shape (N, M) in homogeneous coordinates, where M >= 2. Returns: - A numpy array of shape (N, M-1) in Cartesian coordinates. + An 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}." + f"Input array must have at least two columns, " + f"but got shape {array.shape}." ) - return array[:, :-1] / array[:, -1, np.newaxis] + return array[:, :-1] / array[:, -1:] -def R_to_quat(R): - # 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])) +@tensor_backend_auto +def R_to_quat( + R: Union[Float[Tensor, "n 3 3"], Float[Tensor, "3 3"]] +) -> Union[Float[Tensor, "n 4"], Float[Tensor, "4"]]: + """ + Convert a batch of rotation matrices or a single rotation matrix from + rotation matrix form to quaternion form. + + Args: + R: A tensor containing either a single (3, 3) rotation matrix or a batch + of (n, 3, 3) rotation matrices. + + Returns: + A tensor of quaternions. If the input is (3, 3), the output will be (4,), + and if the input is (n, 3, 3), the output will be (n, 4). + + Ref: + https://github.com/isl-org/StableViewSynthesis/tree/main/co + """ + orig_shape = R.shape + R = ivy.reshape(R, (-1, 3, 3)) + q = create_empty((R.shape[0], 4), dtype=R.dtype, backend=get_tensor_backend(R)) + q[:, 0] = ivy.sqrt(ivy.maximum(0, 1 + R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2])) + q[:, 1] = ivy.sqrt(ivy.maximum(0, 1 + R[:, 0, 0] - R[:, 1, 1] - R[:, 2, 2])) + q[:, 2] = ivy.sqrt(ivy.maximum(0, 1 - R[:, 0, 0] + R[:, 1, 1] - R[:, 2, 2])) + q[:, 3] = ivy.sqrt(ivy.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() + q = q / ivy.vector_norm(q, axis=1, keepdims=True) + + # Handle different input shapes for squeezing + if orig_shape == (3, 3): + return ivy.squeeze(q) + else: + return q def T_to_C(T): @@ -287,7 +307,6 @@ def R_t_to_C(R, t): # 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 diff --git a/camtools/metric.py b/camtools/metric.py index 7cc8bad..9697ac2 100644 --- a/camtools/metric.py +++ b/camtools/metric.py @@ -272,9 +272,12 @@ def _check_inputs( im_mask: np.ndarray = None, ) -> None: # Instance type. - sanity.assert_numpy(im_pd, name="im_pd") - sanity.assert_numpy(im_gt, name="im_gt") - sanity.assert_numpy(im_mask, name="im_mask") + if not isinstance(im_pd, np.ndarray): + raise ValueError(f"im_pd must be numpy array, but got {type(im_pd)}") + if not isinstance(im_gt, np.ndarray): + raise ValueError(f"im_gt must be numpy array, but got {type(im_gt)}") + if im_mask is not None and not isinstance(im_mask, np.ndarray): + raise ValueError(f"im_mask must be numpy array, but got {type(im_mask)}") # Dtype. if im_pd.dtype != np.float32: diff --git a/camtools/sanity.py b/camtools/sanity.py index 4fe66f0..e433b12 100644 --- a/camtools/sanity.py +++ b/camtools/sanity.py @@ -1,12 +1,6 @@ import numpy as np -def assert_numpy(x, name=None): - if not isinstance(x, np.ndarray): - maybe_name = f" {name}" if name is not None else "" - raise ValueError(f"Expected{maybe_name} to be numpy array, but got {type(x)}.") - - def assert_K(K): if K.shape != (3, 3): raise ValueError(f"K must has shape (3, 3), but got {K} of shape {K.shape}.") diff --git a/test/test_convert.py b/test/test_convert.py index 2d28ce0..284f7bd 100644 --- a/test/test_convert.py +++ b/test/test_convert.py @@ -1,10 +1,349 @@ import numpy as np import camtools as ct import pytest +from camtools.backend import is_torch_available, torch +import pytest +import warnings np.set_printoptions(formatter={"float": "{: 0.2f}".format}) +@pytest.fixture(autouse=True) +def ignore_ivy_warnings(): + warnings.filterwarnings( + "ignore", + message=".*Compositional function.*array_mode is set to False.*", + category=UserWarning, + ) + yield + + +def test_pad_0001(): + # Define numpy arrays for testing + in_val_2d = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ], + dtype=np.float64, + ) + gt_out_val_2d = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [0, 0, 0, 1], + ], + dtype=np.float64, + ) + in_val_3d = np.array( + [ + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ], + [ + [13, 14, 15, 16], + [17, 18, 19, 20], + [21, 22, 23, 24], + ], + ], + dtype=np.float64, + ) + gt_out_val_3d = np.array( + [ + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [0, 0, 0, 1], + ], + [ + [13, 14, 15, 16], + [17, 18, 19, 20], + [21, 22, 23, 24], + [0, 0, 0, 1], + ], + ], + dtype=np.float64, + ) + wrong_in_val = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + ], + dtype=np.float64, + ) + + # Test numpy operations + out_val_2d = ct.convert.pad_0001(in_val_2d) + assert isinstance(out_val_2d, np.ndarray) + np.testing.assert_array_equal(out_val_2d, gt_out_val_2d) + + out_val_3d = ct.convert.pad_0001(in_val_3d) + assert isinstance(out_val_3d, np.ndarray) + np.testing.assert_array_equal(out_val_3d, gt_out_val_3d) + + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.pad_0001(wrong_in_val) + + # Test torch operations + if not is_torch_available(): + return + + out_val_2d = ct.convert.pad_0001(torch.from_numpy(in_val_2d)) + assert isinstance(out_val_2d, torch.Tensor) + assert torch.equal(out_val_2d, torch.from_numpy(gt_out_val_2d)) + + out_val_3d = ct.convert.pad_0001(torch.from_numpy(in_val_3d)) + assert isinstance(out_val_3d, torch.Tensor) + assert torch.equal(out_val_3d, torch.from_numpy(gt_out_val_3d)) + + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.pad_0001(torch.from_numpy(wrong_in_val)) + + +def test_rm_pad_0001(): + # Create padded inputs and ground truth outputs for 2D and 3D cases + in_val_2d = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [0, 0, 0, 1], + ], + dtype=np.float64, + ) + gt_out_val_2d = np.array( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ], + dtype=np.float64, + ) + in_val_3d = np.array( + [ + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [0, 0, 0, 1], + ], + [ + [13, 14, 15, 16], + [17, 18, 19, 20], + [21, 22, 23, 24], + [0, 0, 0, 1], + ], + ], + dtype=np.float64, + ) + gt_out_val_3d = np.array( + [ + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ], + [ + [13, 14, 15, 16], + [17, 18, 19, 20], + [21, 22, 23, 24], + ], + ], + dtype=np.float64, + ) + + # Test numpy operations + out_val_2d = ct.convert.rm_pad_0001(in_val_2d) + assert isinstance(out_val_2d, np.ndarray) + np.testing.assert_array_equal(out_val_2d, gt_out_val_2d) + + out_val_3d = ct.convert.rm_pad_0001(in_val_3d) + assert isinstance(out_val_3d, np.ndarray) + np.testing.assert_array_equal(out_val_3d, gt_out_val_3d) + + # Test cases with incorrect bottom row + in_val_2d_wrong = np.copy(in_val_2d) + in_val_2d_wrong[-1, :] = [1, 1, 1, 1] + with pytest.raises(ValueError, match="Expected bottom row to be .*"): + ct.convert.rm_pad_0001(in_val_2d_wrong, check_vals=True) + + # Test torch operations if available + if not is_torch_available(): + return + out_val_2d = ct.convert.rm_pad_0001(torch.from_numpy(in_val_2d)) + assert isinstance(out_val_2d, torch.Tensor) + assert torch.equal(out_val_2d, torch.from_numpy(gt_out_val_2d)) + + out_val_3d = ct.convert.rm_pad_0001(torch.from_numpy(in_val_3d)) + assert isinstance(out_val_3d, torch.Tensor) + assert torch.equal(out_val_3d, torch.from_numpy(gt_out_val_3d)) + + in_val_2d_wrong = torch.from_numpy(in_val_2d_wrong) + with pytest.raises(ValueError, match="Expected bottom row to be .*"): + ct.convert.rm_pad_0001(in_val_2d_wrong, check_vals=True) + + +def test_to_homo(): + # Test with a numpy array + in_val = np.array( + [ + [2, 3], + [4, 5], + [6, 7], + ], + dtype=np.float32, + ) + gt_out_val = np.array( + [ + [2, 3, 1], + [4, 5, 1], + [6, 7, 1], + ], + dtype=np.float32, + ) + out_val = ct.convert.to_homo(in_val) + assert isinstance(out_val, np.ndarray) + np.testing.assert_array_equal(out_val, gt_out_val) + + incorrect_shape_input = np.array([1, 2, 3]) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.to_homo(incorrect_shape_input) + + # Test with a torch tensor + if not is_torch_available(): + return + in_val = torch.from_numpy(in_val) + gt_out_val = torch.from_numpy(gt_out_val) + out_val = ct.convert.to_homo(in_val) + assert isinstance(out_val, torch.Tensor) + assert torch.equal(out_val, gt_out_val) + + incorrect_shape_input = torch.from_numpy(incorrect_shape_input) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.to_homo(incorrect_shape_input) + + +def test_from_homo(): + # Test with a numpy array + in_val = np.array( + [ + [2, 3, 1], + [4, 6, 2], + [6, 9, 3], + ], + dtype=np.float32, + ) + gt_out_val = np.array( + [ + [2, 3], + [2, 3], + [2, 3], + ], + dtype=np.float32, + ) + + # Regular case + out_val = ct.convert.from_homo(in_val) + assert isinstance(out_val, np.ndarray) + np.testing.assert_array_almost_equal(out_val, gt_out_val) + + # Not a 2D array + incorrect_in_val_a = np.array([1, 2, 3], dtype=np.float32) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.from_homo(incorrect_in_val_a) + + # 2D but only one column + incorrect_in_val_b = np.array([[1]], dtype=np.float32) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.from_homo(incorrect_in_val_b) + + # Test with a torch tensor + if not is_torch_available(): + return + in_val = torch.from_numpy(in_val) + gt_out_val = torch.from_numpy(gt_out_val) + + # Regular case + out_val = ct.convert.from_homo(in_val) + assert isinstance(out_val, torch.Tensor) + assert torch.equal(out_val, gt_out_val) + + # Not a 2D array + incorrect_in_val_a = torch.from_numpy(incorrect_in_val_a) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.from_homo(incorrect_in_val_a) + + # 2D but only one column + incorrect_in_val_b = torch.from_numpy(incorrect_in_val_b) + with pytest.raises(ValueError, match=".*got shape.*"): + ct.convert.from_homo(incorrect_in_val_b) + + +def test_R_to_quat(): + theta = np.pi / 2 + R_x = np.array( + [ + [1, 0, 0], + [0, np.cos(theta), -np.sin(theta)], + [0, np.sin(theta), np.cos(theta)], + ] + ) + R_y = np.array( + [ + [np.cos(theta), 0, np.sin(theta)], + [0, 1, 0], + [-np.sin(theta), 0, np.cos(theta)], + ] + ) + R_z = np.array( + [ + [np.cos(theta), -np.sin(theta), 0], + [np.sin(theta), np.cos(theta), 0], + [0, 0, 1], + ] + ) + R_batch = np.array([R_x, R_y, R_z]) + + # Expected quaternions + gt_quat_x = np.array([np.cos(theta / 2), np.sin(theta / 2), 0, 0]) + gt_quat_y = np.array([np.cos(theta / 2), 0, np.sin(theta / 2), 0]) + gt_quat_z = np.array([np.cos(theta / 2), 0, 0, np.sin(theta / 2)]) + gt_quat_batch = np.array([gt_quat_x, gt_quat_y, gt_quat_z]) + + # Test numpy backend + output_x = ct.convert.R_to_quat(R_x) + output_y = ct.convert.R_to_quat(R_y) + output_z = ct.convert.R_to_quat(R_z) + output_batch = ct.convert.R_to_quat(R_batch) + np.testing.assert_allclose(output_x, gt_quat_x, atol=1e-5) + np.testing.assert_allclose(output_y, gt_quat_y, atol=1e-5) + np.testing.assert_allclose(output_z, gt_quat_z, atol=1e-5) + np.testing.assert_allclose(output_batch, gt_quat_batch, atol=1e-5) + + # Test torch backend + if not is_torch_available(): + return + R_x_torch = torch.from_numpy(R_x) + R_y_torch = torch.from_numpy(R_y) + R_z_torch = torch.from_numpy(R_z) + R_batch_torch = torch.from_numpy(R_batch) + output_x_torch = ct.convert.R_to_quat(R_x_torch) + output_y_torch = ct.convert.R_to_quat(R_y_torch) + output_z_torch = ct.convert.R_to_quat(R_z_torch) + output_batch_torch = ct.convert.R_to_quat(R_batch_torch) + assert torch.allclose(output_x_torch, torch.from_numpy(gt_quat_x), atol=1e-5) + assert torch.allclose(output_y_torch, torch.from_numpy(gt_quat_y), atol=1e-5) + assert torch.allclose(output_z_torch, torch.from_numpy(gt_quat_z), atol=1e-5) + assert torch.allclose( + output_batch_torch, torch.from_numpy(gt_quat_batch), atol=1e-5 + ) + + def test_R_t_to_cameracenter(): T = np.array( [ @@ -19,7 +358,10 @@ def test_R_t_to_cameracenter(): expected_camera_center = [-10.7635, 11.8896, 1.348] camera_center = ct.convert.R_t_to_C(R, t) np.testing.assert_allclose( - expected_camera_center, camera_center, rtol=1e-5, atol=1e-5 + expected_camera_center, + camera_center, + rtol=1e-5, + atol=1e-5, ) @@ -235,53 +577,3 @@ def gen_random_T(): rtol=1e-5, atol=1e-5, ) - - -def test_to_homo(): - # Regular case - src = np.array( - [ - [1, 2], - [3, 4], - ] - ) - dst_gt = np.array( - [ - [1, 2, 1], - [3, 4, 1], - ] - ) - dst = ct.convert.to_homo(src) - np.testing.assert_array_equal(dst, dst_gt) - - # Exception case - with pytest.raises(ValueError) as _: - src = np.array([1, 2, 3]) - ct.convert.to_homo(src) - - -def test_from_homo(): - src = np.array( - [ - [2, 4, 2], - [6, 8, 1], - ] - ) - dst_gt = np.array( - [ - [1, 2], - [6, 8], - ] - ) - dst = ct.convert.from_homo(src) - np.testing.assert_array_equal(dst, dst_gt) - - # Exception case for non-2D input - with pytest.raises(ValueError) as _: - src = np.array([1, 2, 3]) - ct.convert.from_homo(src) - - # Exception case for insufficient columns - with pytest.raises(ValueError) as _: - src = np.array([[1]]) - ct.convert.from_homo(src)