From d3dc9c5011c2b73e908f101f84bbe1082db4f8c9 Mon Sep 17 00:00:00 2001 From: Alexander Kristoffersen Date: Fri, 9 Sep 2022 16:10:04 -0700 Subject: [PATCH] Unified place for all pose util methods. (#371) * init pose utils. * more tests. * pose utils applied throughout codebase. * added normalization. --- nerfactory/cameras/camera_paths.py | 7 +-- nerfactory/utils/poses.py | 82 +++++++++++++++++++++++++ tests/utils/test_poses.py | 97 ++++++++++++++++++++++++++++++ 3 files changed, 182 insertions(+), 4 deletions(-) create mode 100644 nerfactory/utils/poses.py create mode 100644 tests/utils/test_poses.py diff --git a/nerfactory/cameras/camera_paths.py b/nerfactory/cameras/camera_paths.py index 776680325..e6f86837a 100644 --- a/nerfactory/cameras/camera_paths.py +++ b/nerfactory/cameras/camera_paths.py @@ -21,6 +21,7 @@ from typing import Optional, Tuple import torch import nerfactory.cameras.utils as camera_utils +import nerfactory.utils.poses as pose_utils from nerfactory.cameras.cameras import Cameras from nerfactory.cameras.utils import get_interpolated_poses_many @@ -82,8 +83,7 @@ def get_spiral_path( target = torch.tensor([0, 0, -focal], device=camera.device) # camera looking in -z direction c2w = camera.camera_to_worlds[0] - ones = torch.tensor([0, 0, 0, 1], device=c2w.device)[None] - c2wh_global = torch.cat([c2w, ones], dim=0) + c2wh_global = pose_utils.to4x4(c2w) local_c2whs = [] for theta in torch.linspace(0.0, 2.0 * torch.pi * rots, steps + 1)[:-1]: @@ -92,8 +92,7 @@ def get_spiral_path( ) lookat = center - target c2w = camera_utils.viewmatrix(lookat, up, center) - ones = torch.tensor([0, 0, 0, 1], device=c2w.device)[None] - c2wh = torch.cat([c2w, ones], dim=0) + c2wh = pose_utils.to4x4(c2w) local_c2whs.append(c2wh) new_c2ws = [] diff --git a/nerfactory/utils/poses.py b/nerfactory/utils/poses.py new file mode 100644 index 000000000..56e7f9d5e --- /dev/null +++ b/nerfactory/utils/poses.py @@ -0,0 +1,82 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Common 3D pose methods +""" + +import torch +from torchtyping import TensorType + + +def to4x4(pose: TensorType[..., 3, 4]) -> TensorType[..., 4, 4]: + """Convert 3x4 pose matrices to a 4x4 with the addition of a homogeneous coordinate. + + Args: + pose (TensorType[..., 3, 4]): Camera pose without homogenous coordinate. + + Returns: + TensorType[..., 4, 4]: Camera poses with additional homogenous coordinate added. + """ + constants = torch.zeros_like(pose[..., :1, :], device=pose.device) + constants[..., :, 3] = 1 + return torch.cat([pose, constants], dim=-2) + + +def inverse(pose: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: + """Invert provided pose matrix. + + Args: + pose (TensorType[..., 3, 4]): Camera pose without homogenous coordinate. + + Returns: + TensorType[..., 3, 4]: Inverse of pose. + """ + R = pose[..., :3, :3] + t = pose[..., :3, 3:] + R_inverse = R.transpose(-2, -1) # pylint: disable=invalid-name + t_inverse = -R_inverse.matmul(t) + return torch.cat([R_inverse, t_inverse], dim=-1) + + +def multiply(pose_a: TensorType[..., 3, 4], pose_b: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: + """Multiply two pose matrices, A @ B. + + Args: + pose_a (TensorType[..., 3, 4]): Left pose matrix, usually a transformation applied to the right. + pose_b (TensorType[..., 3, 4]): Right pose matrix, usually a camera pose that will be tranformed by pose_a. + + Returns: + TensorType[..., 3, 4]: Camera pose matrix where pose_a was applied to pose_b. + """ + R1, t1 = pose_a[..., :3, :3], pose_a[..., :3, 3:] + R2, t2 = pose_b[..., :3, :3], pose_b[..., :3, 3:] + R = R1.matmul(R2) + t = t1 + R1.matmul(t2) + return torch.cat([R, t], dim=-1) + + +def normalize(poses: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: + """Normalize the XYZs of poses to fit within a unit cube ([-1, 1]). Note: This operation is not in-place. + + Args: + poses (TensorType[..., 3, 4]): A collection of poses to be normalized. + + Returns; + TensorType[..., 3, 4]: Normalized collection of poses. + """ + pose_copy = torch.clone(poses) + pose_copy[..., :3, 3] /= torch.max(torch.abs(poses[..., :3, 3])) + + return pose_copy diff --git a/tests/utils/test_poses.py b/tests/utils/test_poses.py new file mode 100644 index 000000000..23b7e192c --- /dev/null +++ b/tests/utils/test_poses.py @@ -0,0 +1,97 @@ +""" +Test pose utils +""" + +import torch + +from nerfactory.utils import poses + + +def test_to4x4(): + """Test addition of homogeneous coordinate to 3D pose.""" + pose = torch.rand((10, 3, 4)) + pose_4x4 = poses.to4x4(pose) + + assert pose_4x4.shape == (10, 4, 4) + assert torch.equal( + pose_4x4[:, 3, 3], + torch.ones( + 10, + ), + ) + + +def test_multiply(): + """Test 3D pose multiplication.""" + pose = torch.tensor( + [ + [ + [1.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 2.0], + [0.0, 0.0, 1.0, 3.0], + ] + ] + ) + translation_pose = poses.multiply(pose, pose) + assert translation_pose.shape == pose.shape + assert torch.equal(translation_pose[..., :, 3], torch.tensor([[2.0, 4.0, 6.0]])) + + pose_a = pose.clone() + pose_a[:, :3, :3] = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ] + ) + pose_b = pose.clone() + pose_b[:, :3, :3] = torch.tensor( + [ + [0.0, -1.0, 0.0], + [0.0, 0.0, -1.0], + [1.0, 0.0, 0.0], + ] + ) + + translation_rotation_pose = poses.multiply(pose_a, pose_b) + assert torch.allclose(translation_rotation_pose, (poses.to4x4(pose_a) @ poses.to4x4(pose_b))[:, :3, :4]) + + +def test_inverse(): + """Test 3D pose inversion.""" + + pose = torch.rand((10, 3, 4)) + pose[:, :3, :3] = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ] + ) + pose_inv = poses.inverse(pose) + + assert pose_inv.shape == pose.shape + + unit_pose = torch.zeros_like(pose) + unit_pose[:, :3, :3] = torch.eye(3) + + assert torch.allclose( + poses.multiply(pose, pose_inv), + unit_pose, + ) + + +def test_normalize(): + """Test 3D pose normalization""" + pose = torch.ones((2, 3, 4)) + pose[:, :3, :3] = torch.eye(3) + pose[0, :3, 3] = torch.tensor([2.0, 0.0, -2.0]) + pose[1, :3, 3] = torch.tensor([1.0, 1.0, 1.0]) + + pose_scaled = poses.normalize(pose) + + assert pose_scaled.shape == pose.shape + assert torch.max(pose_scaled[:, :3, 3]) <= 1.0 + + assert torch.equal(pose_scaled[0, :3, 3], torch.tensor([1.0, 0.0, -1.0])) + assert torch.equal(pose_scaled[1, :3, 3], torch.tensor([0.5, 0.5, 0.5]))