From 3841e2b3950d54865615b60b4ca79231ffa754a6 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Tue, 6 Sep 2022 00:56:47 +0100 Subject: [PATCH] Auto orient poses (#342) --- nerfactory/cameras/utils.py | 77 +++++++++++++++++++++++++- nerfactory/configs/base.py | 3 +- nerfactory/dataloaders/data_parsers.py | 17 +++--- 3 files changed, 88 insertions(+), 9 deletions(-) diff --git a/nerfactory/cameras/utils.py b/nerfactory/cameras/utils.py index 5be916049..2cd458b3a 100644 --- a/nerfactory/cameras/utils.py +++ b/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 diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 3baaed5bc..11a4ae65a 100644 --- a/nerfactory/configs/base.py +++ b/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 diff --git a/nerfactory/dataloaders/data_parsers.py b/nerfactory/dataloaders/data_parsers.py index 0c4b895be..4a5b80c0e 100644 --- a/nerfactory/dataloaders/data_parsers.py +++ b/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, )