diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 6fd85b358..5ae29ca49 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/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 diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 1e29f6b11..085904d63 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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. diff --git a/nerfstudio/data/dataparsers/record3d_dataparser.py b/nerfstudio/data/dataparsers/record3d_dataparser.py index ef20f21ba..e3ba5d858 100644 --- a/nerfstudio/data/dataparsers/record3d_dataparser.py +++ b/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)