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

@@ -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 = []

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

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]))