add camera converter for gsplats (#813)

* add camera converter

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>

* fix few tests tolerance / flakiness

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>

---------

Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>

Author: Caenorst

docs/kaolin_ext.py

@@ -93,6 +93,7 @@ def run_apidoc(_):
             "kaolin/render/camera/intrinsics.py",
             "kaolin/render/camera/legacy.py",
             "kaolin/render/camera/raygen.py",
+            "kaolin/render/camera/gsplats.py",
             "kaolin/non_commercial/flexicubes/",
             "kaolin/non_commercial/flexicubes/flexicubes.py",
             "kaolin/non_commercial/flexicubes/tables.py"

kaolin/render/camera/__init__.py

@@ -21,5 +21,6 @@
 from .coordinates import *
 from .legacy import *
 from .raygen import *
+from .gsplats import *
 
 __all__ = [k for k in locals().keys() if not k.startswith('__')]

kaolin/render/camera/gsplats.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+from .camera import Camera
+from .intrinsics import CameraFOV
+
+__all__ = [
+    'kaolin_camera_to_gsplats',
+    'gsplats_camera_to_kaolin'
+]
+
+
+
+def kaolin_camera_to_gsplats(kal_camera, gs_cam_cls):
+    """Convert Kaolin Camera to `INRIA gaussian splats`_ camera.
+
+    .. note::
+       This has been tested with the version commit `472689c`_
+
+    Args:
+        kal_camera (Camera): camera to convert.
+        gs_cam_cls (class):
+            This is the gsplats ``Camera`` class,
+            usually located in gsplats/scene/cameras.py.
+
+    Returns:
+        (gsplats.scene.cameras.Camera): converted INRIA gsplats camera.
+
+    .. _INRIA gaussian splats:
+        https://github.com/graphdeco-inria/gaussian-splatting
+    .. _472689c:
+        https://github.com/graphdeco-inria/gaussian-splatting/tree/472689c0dc70417448fb451bf529ae532d32c095
+    """
+    R = kal_camera.extrinsics.R[0].clone()
+    R[1:3] = -R[1:3]
+    T = kal_camera.extrinsics.t.squeeze()
+    T[1:3] = -T[1:3]
+    return gs_cam_cls(colmap_id=0,
+                      R=R.transpose(1, 0).cpu().numpy(),
+                      T=T.cpu().numpy(),
+                      FoVx=kal_camera.fov(CameraFOV.HORIZONTAL, in_degrees=False),
+                      FoVy=kal_camera.fov(CameraFOV.VERTICAL, in_degrees=False),
+                      image=torch.zeros((3, kal_camera.height, kal_camera.width)),  # fake
+                      gt_alpha_mask=None,
+                      image_name='fake',
+                      uid=0)
+
+def gsplats_camera_to_kaolin(gs_camera):
+    """Convert `INRIA gaussian splats`_ camera to Kaolin camera.
+
+    .. note::
+       This has been tested with the version commit `472689c`_
+
+    Args:
+        gs_camera (gsplats.scene.cameras.Camera): camera to convert.
+
+    Returns:
+        (Camera): converted Kaolin camera.
+
+    .. _INRIA gaussian splats:
+        https://github.com/graphdeco-inria/gaussian-splatting
+    .. _472689c:
+        https://github.com/graphdeco-inria/gaussian-splatting/tree/472689c0dc70417448fb451bf529ae532d32c095
+    """
+    view_mat = gs_camera.world_view_transform.transpose(1, 0).clone()
+    view_mat[1:3] = -view_mat[1:3]
+    res = Camera.from_args(
+        view_matrix=view_mat,
+        width=gs_camera.image_width, height=gs_camera.image_height,
+        fov=gs_camera.FoVy, device='cpu')
+    return res

kaolin/render/lighting/sg.py

@@ -148,7 +148,7 @@ def sg_from_sun(direction, strength, angle, color):
 
     Args:
         strength (torch.Tensor):
-            The strength of the suns, of shape :math:`(\text{num_suns},)`, [0..inf] expected,
+            The strength of the suns, of shape :math:`(\text{num_suns},)`, [1..inf] expected,
             usually in low integer range.
         color (torch.Tensor):
             The color of the suns,

tests/python/kaolin/ops/mesh/test_mesh.py

@@ -170,7 +170,7 @@ def test_average_face_vertex_features(self, device, dtype):
 
         vertex_features = mesh.average_face_vertex_features(faces, face_features)
         if dtype == torch.half:
-            assert torch.allclose(expected, vertex_features, rtol=1e-03, atol=1e-05)
+            assert torch.allclose(expected, vertex_features, rtol=1e-3, atol=1e-3)
         else:
             assert torch.allclose(expected, vertex_features)
 

tests/python/kaolin/ops/test_pointcloud.py

@@ -35,6 +35,7 @@ def test_center_points(device, dtype):
     points[:, 0, :] = 1.0  # make sure 1 is included
     points[:, 1, :] = 0.0  # make sure 0 is included
     points = points - 0.5  # -0.5...0.5
+    points = torch.clamp((torch.sign(points) * 1e-3) + points, -0.5, 0.5)
 
     factors = 0.2 + 2 * torch.rand((B, 1, 1), device=device, dtype=dtype)
     translations = torch.rand((B, 1, 3), device=device, dtype=dtype) - 0.5
@@ -66,4 +67,4 @@ def test_center_points(device, dtype):
     points[0, :, :] = torch.tensor([0, 2., 4.], dtype=dtype, device=device).reshape((1, 3))
     points_centered = kaolin.ops.pointcloud.center_points(points * factors + translations, normalize=True)
     points[0, :, :] = 0
-    assert torch.allclose(points, points_centered, atol=atol, rtol=rtol)
+    assert torch.allclose(points, points_centered, atol=atol, rtol=rtol)

tests/python/kaolin/render/camera/test_gsplats.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import sys
+import subprocess
+import pytest
+import math
+import random
+import torch
+import shutil
+import numpy as np
+from git import Repo
+from kaolin.render.camera import kaolin_camera_to_gsplats, gsplats_camera_to_kaolin, Camera
+
+# dealing with nvcr
+if torch.version.cuda == '12.5':
+    pytest.skip("gsplats is not installable with CUDA 12.5", allow_module_level=True)
+
+ROOT_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'gsplats')
+@pytest.fixture(scope="module")
+def gs_cam_cls():
+    repo = Repo.clone_from(
+        url='https://github.com/graphdeco-inria/gaussian-splatting',
+        multi_options=['--recursive'],
+        to_path=ROOT_DIR
+    )
+    sys.path.append(ROOT_DIR)
+    subprocess.check_call([
+        sys.executable, "-m", "pip", "install",
+        os.path.join(ROOT_DIR, "submodules", "diff-gaussian-rasterization")
+    ])
+    subprocess.check_call([
+        sys.executable, "-m", "pip", "install",
+        os.path.join(ROOT_DIR, "submodules", "simple-knn")
+    ])
+    from .gsplats.scene.cameras import Camera as GSCamera
+    
+    yield GSCamera
+    sys.path.remove(ROOT_DIR)
+    shutil.rmtree(ROOT_DIR)
+
+    
+
+class TestGsplats:
+    def test_cycle(self, gs_cam_cls):
+        kal_cam = Camera.from_args(
+            eye=torch.rand((3,)),
+            at=torch.rand((3,)),
+            up=torch.nn.functional.normalize(torch.rand((3,)), dim=0),
+            fov=random.random() * math.pi,
+            width=512, height=512,
+        )
+        gs_cam = kaolin_camera_to_gsplats(kal_cam, gs_cam_cls)
+        out_cam = gsplats_camera_to_kaolin(gs_cam)
+        assert torch.allclose(out_cam, kal_cam)
+    
+    def test_kaolin_to_gsplats_regression(self, gs_cam_cls):
+        kal_cam = Camera.from_args(
+            eye=torch.tensor([1., 2., 3.]), 
+            at=torch.tensor([0.3, 0.1, 0.2]),
+            up=torch.tensor([0., 1., 0.]),
+            fov=math.pi / 4,
+            width=512, height=512,
+        )
+        gs_cam = kaolin_camera_to_gsplats(kal_cam, gs_cam_cls)
+        expected_R = np.array([[ 0.9701425,   0.13336042, -0.20257968],
+                               [ 0.,         -0.83525735, -0.5498591 ],
+                               [-0.24253562,  0.53344166, -0.8103187 ]])
+        expected_T = np.array([-0.24253559, -0.06317067,  3.733254  ])
+        expected_fovx = torch.tensor([0.7854])
+        expected_fovy = torch.tensor([0.7854])
+        assert np.allclose(expected_R, gs_cam.R)
+        assert np.allclose(expected_T, gs_cam.T)
+        assert torch.allclose(expected_fovx, gs_cam.FoVx)
+        assert torch.allclose(expected_fovy, gs_cam.FoVy)

tests/python/kaolin/render/lighting/test_sg.py

@@ -420,7 +420,7 @@ def test_from_sun(self):
 
 class TestUtilities:
     def test_sg_from_sun(self):
-        strength = torch.rand((7,), dtype=torch.float) * 10
+        strength = torch.rand((7,), dtype=torch.float) * 10 + 1.
         direction = torch.rand((7, 3), dtype=torch.float)
         angle = torch.rand((7,), dtype=torch.float) * math.pi
         color = torch.rand((7, 3), dtype=torch.float)

tools/ci_requirements.txt

@@ -7,3 +7,5 @@ pytest==7.3.0
 pytest-cov==3.0.0
 nbmake==1.4.1
 jupyter
+plyfile
+gitpython