diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 89c7c4e21..12367e51f 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -24,6 +24,7 @@ from typing import Any, Dict, List, Optional, Type import torch from torchtyping import TensorType +from typing_extensions import Literal import nerfstudio.configs.base_config as cfg from nerfstudio.cameras.cameras import Cameras @@ -89,6 +90,26 @@ class DataparserOutputs: with open(path, "w", encoding="UTF-8") as file: json.dump(data, file, indent=4) + def transform_poses_to_original_space( + self, + poses: TensorType["num_poses", 3, 4], + camera_convention: Literal["opengl", "opencv"] = "opencv", + ) -> TensorType["num_poses", 3, 4]: + """ + Transforms the poses in the transformed space back to the original world coordinate system. + Args: + poses: Poses in the transformed space + camera_convention: Camera system convention used for the transformed poses + Returns: + Original poses + """ + return transform_poses_to_original_space( + poses, + self.dataparser_transform, + self.dataparser_scale, + camera_convention=camera_convention, + ) + @dataclass class DataParserConfig(cfg.InstantiateConfig): @@ -136,3 +157,46 @@ class DataParser: """ dataparser_outputs = self._generate_dataparser_outputs(split) return dataparser_outputs + + +def transform_poses_to_original_space( + poses: TensorType["num_poses", 3, 4], + applied_transform: TensorType[3, 4], + applied_scale: float, + camera_convention: Literal["opengl", "opencv"] = "opencv", +) -> TensorType["num_poses", 3, 4]: + """ + Transforms the poses in the transformed space back to the original world coordinate system. + Args: + poses: Poses in the transformed space + applied_transform: Transform matrix applied in the data processing step + applied_scale: Scale used in the data processing step + camera_convention: Camera system convention used for the transformed poses + Returns: + Original poses + """ + output_poses = torch.cat( + ( + poses, + torch.tensor([[[0, 0, 0, 1]]], dtype=poses.dtype, device=poses.device).repeat_interleave(len(poses), 0), + ), + 1, + ) + output_poses[..., :3, 3] /= applied_scale + inv_transform = torch.linalg.inv( + torch.cat( + ( + applied_transform, + torch.tensor([[0, 0, 0, 1]], dtype=applied_transform.dtype, device=applied_transform.device), + ), + 0, + ) + ) + output_poses = torch.einsum("ij,bjk->bik", inv_transform, output_poses) + if camera_convention == "opencv": + output_poses[..., 0:3, 1:3] *= -1 + elif camera_convention == "opengl": + pass + else: + raise ValueError(f"Camera convention {camera_convention} is not supported.") + return output_poses[:, :3] diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index 3e4fa82a9..1b86a5adb 100644 --- a/tests/process_data/test_process_images.py +++ b/tests/process_data/test_process_images.py @@ -12,12 +12,30 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars from nerfstudio.data.utils.colmap_parsing_utils import Camera from nerfstudio.data.utils.colmap_parsing_utils import Image as ColmapImage from nerfstudio.data.utils.colmap_parsing_utils import ( + qvec2rotmat, write_cameras_binary, write_images_binary, ) from scripts.process_data import ProcessImages +def random_quaternion(num_poses: int): + """ + Generates random rotation quatenion. + """ + u, v, w = np.random.uniform(size=(3, num_poses)) + quaternion = np.stack( + ( + np.sqrt(1 - u) * np.sin(2 * np.pi * v), + np.sqrt(1 - u) * np.cos(2 * np.pi * v), + np.sqrt(u) * np.sin(2 * np.pi * w), + np.sqrt(u) * np.cos(2 * np.pi * w), + ), + -1, + ) + return quaternion + + def test_process_images_skip_colmap(tmp_path: Path): """ Test ns-process-data images @@ -34,21 +52,21 @@ def test_process_images_skip_colmap(tmp_path: Path): ) frames = {} num_frames = 10 - qvecs = np.random.uniform(size=(num_frames, 4)) + qvecs = random_quaternion(num_frames) tvecs = np.random.uniform(size=(num_frames, 3)) - # original_poses = np.concatenate( - # ( - # np.concatenate( - # ( - # np.stack(list(map(qvec2rotmat, qvecs))), - # tvecs[:, :, None], - # ), - # -1, - # ), - # np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0), - # ), - # -2, - # ) + original_poses = np.concatenate( + ( + np.concatenate( + ( + np.stack(list(map(qvec2rotmat, qvecs))), + tvecs[:, :, None], + ), + -1, + ), + np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0), + ), + -2, + ) for i in range(num_frames): frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"image_{i}.png", [], []) Image.new("RGB", (width, height)).save(tmp_path / "images" / f"image_{i}.png") @@ -77,4 +95,12 @@ def test_process_images_skip_colmap(tmp_path: Path): outputs = parser.get_dataparser_outputs("train") assert len(outputs.image_filenames) == 9 assert torch.is_tensor(outputs.dataparser_transform) - # TODO @jkulhanek: Add tests if the loaded poses are the same + + # Test if the original poses can be obtained back + gt_poses = original_poses[[0, 1, 2, 3, 4, 5, 6, 7, 9]] + dataparser_poses = outputs.transform_poses_to_original_space(outputs.cameras.camera_to_worlds, "opencv").numpy() + dataparser_poses = np.concatenate( + (dataparser_poses, np.array([[[0, 0, 0, 1]]]).repeat(len(dataparser_poses), 0)), 1 + ) + dataparser_poses = np.linalg.inv(dataparser_poses) + np.testing.assert_allclose(gt_poses, dataparser_poses, rtol=0, atol=1e-5)