Adds util function to convert transformed poses back + tests (#1594)

Improves tests of ns-process-data images +
Adds util function to convert poses back from transformed coordinates.

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]

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)