Add more controls for orienting poses (#800)

nerfstudio/cameras/camera_utils.py

@@ -406,20 +406,21 @@ def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
     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"
+def auto_orient_and_center_poses(
+    poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True
 ) -> 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".
+        method: The method to use for orientation.
+        center_poses: If True, the poses are centered around the origin.
 
     Returns:
         The oriented poses.
@@ -428,16 +429,21 @@ def auto_orient_poses(
     translation = poses[..., :3, 3]
 
     mean_translation = torch.mean(translation, dim=0)
-    translation = translation - mean_translation
+    translation_diff = translation - mean_translation
+
+    if center_poses:
+        translation = mean_translation
+    else:
+        translation = torch.zeros_like(mean_translation)
 
     if method == "pca":
-        _, eigvec = torch.linalg.eigh(translation.T @ translation)
+        _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff)
         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)
+        transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1)
         oriented_poses = transform @ poses
 
         if oriented_poses.mean(axis=0)[2, 1] < 0:
@@ -447,7 +453,10 @@ def auto_orient_poses(
         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)
+        transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1)
         oriented_poses = transform @ poses
+    elif method == "none":
+        oriented_poses = poses
+        poses[:, 3, :3] -= translation
 
     return oriented_poses

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -53,8 +53,12 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px."""
     scene_scale: float = 1.0
     """How much to scale the region of interest by."""
-    orientation_method: Literal["pca", "up"] = "up"
+    orientation_method: Literal["pca", "up", "none"] = "up"
     """The method to use for orientation."""
+    center_poses: bool = True
+    """Whether to center the poses."""
+    auto_scale_poses: bool = True
+    """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
     train_split_percentage: float = 0.9
     """The percent of images to use for training. The remaining images are for eval."""
 
@@ -112,10 +116,15 @@ class Nerfstudio(DataParser):
             raise ValueError(f"Unknown dataparser split {split}")
 
         poses = torch.from_numpy(np.array(poses).astype(np.float32))
-        poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method)
+        poses = camera_utils.auto_orient_and_center_poses(
+            poses, method=self.config.orientation_method, center_poses=self.config.center_poses
+        )
 
         # Scale poses
-        scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3]))
+        scale_factor = 1.0
+        if self.config.auto_scale_poses:
+            scale_factor /= torch.max(torch.abs(poses[:, :3, 3]))
+
         poses[:, :3, 3] *= scale_factor * self.config.scale_factor
 
         # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.

nerfstudio/data/dataparsers/record3d_dataparser.py

@@ -105,9 +105,9 @@ class Record3D(DataParser):
         # convert to Tensors
         poses = torch.from_numpy(poses[:, :3, :4])
 
-        poses = camera_utils.auto_orient_poses(pose_utils.to4x4(poses), method=self.config.orientation_method)[
-            :, :3, :4
-        ]
+        poses = camera_utils.auto_orient_and_center_poses(
+            pose_utils.to4x4(poses), method=self.config.orientation_method
+        )[:, :3, :4]
 
         # Centering poses
         poses[:, :3, 3] = poses[:, :3, 3] - torch.mean(poses[:, :3, 3], dim=0)