Auto orient poses (#342)

nerfactory/cameras/utils.py

@@ -18,7 +18,7 @@ Camera transformation helper code.
 """
 
 import math
-from typing import Tuple
+from typing import Literal, Tuple
 
 import numpy as np
 import torch
@@ -459,3 +459,78 @@ def _compute_residual_and_jacobian(
     fy_y = d + d_y * y + 2.0 * p2 * x + 6.0 * p1 * y
 
     return fx, fy, fx_x, fx_y, fy_x, fy_y
+
+
+def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
+    """Compute the rotation matrix that rotates vector a to vector b.
+
+    Args:
+        a: The vector to rotate.
+        b: The vector to rotate to.
+    Returns:
+        The rotation matrix.
+    """
+    a = a / torch.linalg.norm(a)
+    b = b / torch.linalg.norm(b)
+    v = torch.cross(a, b)
+    c = torch.dot(a, b)
+    # If vectors are exactly opposite, we add a little noise to one of them
+    if c < -1 + 1e-8:
+        eps = (torch.rand(3) - 0.5) * 0.01
+        return rotmat(a + eps, b)
+    s = torch.linalg.norm(v)
+    skew_sym_mat = torch.Tensor(
+        [
+            [0, -v[2], v[1]],
+            [v[2], 0, -v[0]],
+            [-v[1], v[0], 0],
+        ]
+    )
+    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
+
+
+def auto_orient_poses(
+    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up"] = "up"
+) -> TensorType["num_poses":..., 3, 4]:
+    """Orients and centers the poses. We provide two methods for orientation: pca and up.
+
+    pca: Orient the poses so that the principal component of the points is aligned with the axes.
+        This method works well when all of the cameras are in the same plane.
+
+    up: Orient the poses so that the average up vector is aligned with the z axis.
+        This method works well when images are not at arbitrary angles.
+
+    Args:
+        poses: The poses to orient.
+        method: The method to use for orientation. Either "pca" or "up".
+
+    Returns:
+        The oriented poses.
+    """
+
+    translation = poses[..., :3, 3]
+
+    mean_translation = torch.mean(translation, dim=0)
+    translation = translation - mean_translation
+
+    if method == "pca":
+        _, eigvec = torch.linalg.eigh(translation.T @ translation)
+        eigvec = torch.flip(eigvec, dims=(-1,))
+
+        if torch.linalg.det(eigvec) < 0:
+            eigvec[:, 2] = -eigvec[:, 2]
+
+        transform = torch.cat([eigvec, eigvec @ -mean_translation[..., None]], dim=-1)
+        oriented_poses = transform @ poses
+
+        if oriented_poses.mean(axis=0)[2, 1] < 0:
+            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
+    elif method == "up":
+        up = torch.mean(poses[:, :3, 1], dim=0)
+        up = up / torch.linalg.norm(up)
+
+        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
+        transform = torch.cat([rotation, rotation @ -mean_translation[..., None]], dim=-1)
+        oriented_poses = transform @ poses
+
+    return oriented_poses

nerfactory/configs/base.py

@@ -178,7 +178,8 @@ class NerfactoryDataParserConfig(DataParserConfig):
     data_directory: Path = Path("data/ours/posterv2")
     scale_factor: float = 1.0
     downscale_factor: int = 1
-    scene_scale: float = 0.33
+    scene_scale: float = 4.0
+    orientation_method: Literal["pca", "up"] = "up"
 
 
 @dataclass

nerfactory/dataloaders/data_parsers.py

@@ -136,8 +136,9 @@ class Nerfactory(DataParser):
     Args:
         data_directory: Location of data
         scale_factor: How much to scale the camera origins by.
-        downscale_factor: How much to downscale images. Defaults to 1.
-        scene_scale: How much to scale the scene. Defaults to 0.33
+        downscale_factor: How much to downscale images.
+        scene_scale: How much to scale the scene.
+        orientation_method: The method to use for orientation. Either "pca" or "up".
     """
 
     config: cfg.NerfactoryDataParserConfig
@@ -169,10 +170,12 @@ class Nerfactory(DataParser):
         No image files found. 
         You should check the file_paths in the transforms.json file to make sure they are correct.
         """
-        poses = np.array(poses).astype(np.float32)
-        poses[:3, 3] *= self.config.scene_scale
+        poses = torch.from_numpy(np.array(poses).astype(np.float32))
+        poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method)
 
-        camera_to_world = torch.from_numpy(poses[:, :3])  # camera to world transform
+        # Scale poses
+        scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3]))
+        poses[:, :3, 3] *= scale_factor
 
         distortion_params = camera_utils.get_distortion_params(
             k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"])
@@ -180,7 +183,7 @@ class Nerfactory(DataParser):
 
         # in x,y,z order
         # assumes that the scene is centered at the origin
-        aabb_scale = 16
+        aabb_scale = self.config.scene_scale
         scene_bounds = SceneBounds(
             aabb=torch.tensor(
                 [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
@@ -195,7 +198,7 @@ class Nerfactory(DataParser):
             distortion_params=distortion_params,
             height=int(meta["h"]),
             width=int(meta["w"]),
-            camera_to_worlds=camera_to_world,
+            camera_to_worlds=poses[:, :3, :4],
             camera_type=CameraType.PERSPECTIVE,
         )