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.

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)