GaussianSplats export: allow uint8 in ply (#3111)

* GaussianSplats export: allow uint8 in ply

* Ruff formatting

* Add test and narrower typing

* Formatting

* Revert narrow typing, not supported by old python

* Fix typing for old pythons

nerfstudio/scripts/exporter.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 import json
 import os
 import sys
+import typing
 from collections import OrderedDict
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -487,15 +488,19 @@ class ExportGaussianSplat(Exporter):
     """Scale of the oriented bounding box along each axis."""
 
     @staticmethod
-    def write_ply(filename: str, count: int, map_to_tensors: OrderedDict[str, np.ndarray]):
+    def write_ply(
+        filename: str,
+        count: int,
+        map_to_tensors: typing.OrderedDict[str, np.ndarray],
+    ):
         """
-        Writes a PLY file with given vertex properties and their float values in the order specified by the OrderedDict.
+        Writes a PLY file with given vertex properties and a tensor of float or uint8 values in the order specified by the OrderedDict.
         Note: All float values will be converted to float32 for writing.
 
         Parameters:
         filename (str): The name of the file to write.
         count (int): The number of vertices to write.
-        map_to_tensors (OrderedDict[str, np.ndarray]): An ordered dictionary mapping property names to numpy arrays of float values.
+        map_to_tensors (OrderedDict[str, np.ndarray]): An ordered dictionary mapping property names to numpy arrays of float or uint8 values.
             Each array should be 1-dimensional and of equal length matching 'count'. Arrays should not be empty.
         """
 
@@ -503,12 +508,14 @@ class ExportGaussianSplat(Exporter):
         if not all(len(tensor) == count for tensor in map_to_tensors.values()):
             raise ValueError("Count does not match the length of all tensors")
 
-        # Type check for numpy arrays of type float and non-empty
+        # Type check for numpy arrays of type float or uint8 and non-empty
         if not all(
-            isinstance(tensor, np.ndarray) and tensor.dtype.kind in ["f", "d"] and tensor.size > 0
+            isinstance(tensor, np.ndarray)
+            and (tensor.dtype.kind == "f" or tensor.dtype == np.uint8)
+            and tensor.size > 0
             for tensor in map_to_tensors.values()
         ):
-            raise ValueError("All tensors must be numpy arrays of float type and not empty")
+            raise ValueError("All tensors must be numpy arrays of float or uint8 type and not empty")
 
         with open(filename, "wb") as ply_file:
             # Write PLY header
@@ -518,17 +525,21 @@ class ExportGaussianSplat(Exporter):
             ply_file.write(f"element vertex {count}\n".encode())
 
             # Write properties, in order due to OrderedDict
-            for key in map_to_tensors.keys():
-                ply_file.write(f"property float {key}\n".encode())
+            for key, tensor in map_to_tensors.items():
+                data_type = "float" if tensor.dtype.kind == "f" else "uchar"
+                ply_file.write(f"property {data_type} {key}\n".encode())
 
             ply_file.write(b"end_header\n")
 
             # Write binary data
-            # Note: If this is a perfromance bottleneck consider using numpy.hstack for efficiency improvement
+            # Note: If this is a performance bottleneck consider using numpy.hstack for efficiency improvement
             for i in range(count):
                 for tensor in map_to_tensors.values():
                     value = tensor[i]
-                    ply_file.write(np.float32(value).tobytes())
+                    if tensor.dtype.kind == "f":
+                        ply_file.write(np.float32(value).tobytes())
+                    elif tensor.dtype == np.uint8:
+                        ply_file.write(value.tobytes())
 
     def main(self) -> None:
         if not self.output_dir.exists():

tests/scripts/test_exporter.py

@@ -0,0 +1,82 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. 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 collections import OrderedDict
+from pathlib import Path
+
+import numpy as np
+import open3d as o3d
+import pytest
+
+from nerfstudio.scripts.exporter import ExportGaussianSplat
+
+
+def test_export_gaussian_splat_write_ply(tmp_path: Path):
+    filename = str(tmp_path / "test_export_gaussian_splat_write_ply.ply")
+    count = 10
+    colors_original = np.random.randint(0, 255, size=(count,), dtype=np.uint8)
+    x_original = np.random.rand(count).astype(np.float32)
+    y_original = np.random.rand(count).astype(np.float32)
+    z_original = np.random.rand(count).astype(np.float64)  # should convert to float32
+
+    map_to_tensors: OrderedDict[str, np.ndarray] = OrderedDict(
+        [
+            ("colors", colors_original),
+            ("x", x_original),
+            ("y", y_original),
+            ("z", z_original),
+        ]
+    )
+
+    ExportGaussianSplat.write_ply(filename, count, map_to_tensors)
+
+    # Read the PLY file
+    pcd = o3d.t.io.read_point_cloud(filename)
+    positions = pcd.point.positions.numpy()
+    colors = pcd.point.colors.numpy()
+
+    # Assertions
+    # Check if the file exists and has the correct number of points
+    assert len(positions) == count, "Mismatch in the number of points written and read."
+
+    # Check colors
+    np.testing.assert_array_equal(colors_original, colors[:, 0], "Colors do not match.")
+
+    # Check x, y, and z coordinates with epsilon tollerence
+    epsilon = 1e-5  # Adjust based on expected precision
+    np.testing.assert_allclose(x_original, positions[:, 0], atol=epsilon, err_msg="X coordinates do not match.")
+    np.testing.assert_allclose(y_original, positions[:, 1], atol=epsilon, err_msg="Y coordinates do not match.")
+    np.testing.assert_allclose(
+        z_original.astype(np.float32), positions[:, 2], atol=epsilon, err_msg="Z coordinates do not match."
+    )
+
+
+def test_export_gaussian_splat_write_ply_mismatched_count(tmp_path: Path):
+    filename = str(tmp_path / "test_export_gaussian_splat_write_ply_mismatched_count.ply")
+    count = 10
+    map_to_tensors: OrderedDict[str, np.ndarray] = OrderedDict(
+        [
+            ("x", np.random.rand(count + 1).astype(np.float32)),
+            ("y", np.random.rand(count).astype(np.float32)),
+            ("z", np.random.rand(count - 1).astype(np.float32)),
+        ]
+    )
+    with pytest.raises(ValueError):
+        ExportGaussianSplat.write_ply(filename, count, map_to_tensors)
+
+
+if __name__ == "__main__":
+    # Run the test
+    test_export_gaussian_splat_write_ply(Path("."))
+    test_export_gaussian_splat_write_ply_mismatched_count(Path("."))