From bc9328c7ff70045fce21838122f48ab5201c4ae3 Mon Sep 17 00:00:00 2001 From: Mathijs Henquet Date: Mon, 29 Apr 2024 20:30:25 +0200 Subject: [PATCH] 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 | 31 ++++++++----- tests/scripts/test_exporter.py | 82 ++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 10 deletions(-) create mode 100644 tests/scripts/test_exporter.py diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index e3f40aab1..7d90e0708 100644 --- a/nerfstudio/scripts/exporter.py +++ b/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(): diff --git a/tests/scripts/test_exporter.py b/tests/scripts/test_exporter.py new file mode 100644 index 000000000..e25e4652e --- /dev/null +++ b/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("."))