diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index 475cc5c38..0820b94f0 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -294,6 +294,26 @@ class SplatfactoModel(Model): raise ValueError(f"""Splatfacto does not support strategy {self.config.strategy} Currently, the supported strategies include default and mcmc.""") + # Surface the rasterization mode (classic vs antialiased) in the viewer. + # The value is baked into the trained model via `rasterize_mode`, so the + # control is disabled/frozen: it reflects how the model was trained (and + # what a PLY exported from it records as its `SplatRenderMode` comment) + # rather than being an interactive toggle. Imported locally to keep the + # viewer dependency out of headless import paths. + from nerfstudio.viewer.viewer_elements import ViewerDropdown + + self.rasterize_mode_viewer = ViewerDropdown( + name="Anti-Aliasing", + default_value="antialiased" if self.config.rasterize_mode == "antialiased" else "classic", + options=["classic", "antialiased"], + disabled=True, + hint=( + "Rasterization mode this model was trained with (frozen). " + "'antialiased' corresponds to a 'SplatRenderMode: mip' comment " + "in exported PLY files; 'classic' to 'SplatRenderMode: default'." + ), + ) + @property def colors(self): if self.config.sh_degree > 0: diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index c26797a5c..89fe66d77 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -503,6 +503,7 @@ class ExportGaussianSplat(Exporter): filename: str, count: int, map_to_tensors: typing.OrderedDict[str, np.ndarray], + antialiased: typing.Optional[bool] = None, ): """ Writes a PLY file with given vertex properties and a tensor of float or uint8 values in the order specified by the OrderedDict. @@ -513,6 +514,11 @@ class ExportGaussianSplat(Exporter): 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 or uint8 values. Each array should be 1-dimensional and of equal length matching 'count'. Arrays should not be empty. + antialiased (Optional[bool]): Whether the model was trained with the + antialiased (mip-splatting) rasterize mode. When not None, a + 'SplatRenderMode' header comment is written ('mip' when antialiased, + 'default' otherwise) so viewers can select a matching rasterization + mode. When None (default), no such comment is written. """ # Ensure count matches the length of all tensors @@ -535,6 +541,11 @@ class ExportGaussianSplat(Exporter): ply_file.write(b"format binary_little_endian 1.0\n") ply_file.write(f"comment Generated by Nerstudio {nerfstudio_version}\n".encode()) ply_file.write(b"comment Vertical Axis: z\n") + # Record the rasterization mode so viewers can match it. 'mip' is the + # antialiased mode; 'default' is classic. See SplatfactoModelConfig. + if antialiased is not None: + render_mode = "mip" if antialiased else "default" + ply_file.write(f"comment SplatRenderMode: {render_mode}\n".encode()) ply_file.write(f"element vertex {count}\n".encode()) # Write properties, in order due to OrderedDict @@ -646,7 +657,10 @@ class ExportGaussianSplat(Exporter): map_to_tensors[k] = map_to_tensors[k][select] count = np.sum(select) - ExportGaussianSplat.write_ply(str(filename), count, map_to_tensors) + # Record whether the model uses the antialiased (mip-splatting) rasterize + # mode so the exported PLY carries a matching SplatRenderMode comment. + antialiased = model.config.rasterize_mode == "antialiased" + ExportGaussianSplat.write_ply(str(filename), count, map_to_tensors, antialiased=antialiased) Commands = tyro.conf.FlagConversionOff[ diff --git a/tests/scripts/test_exporter.py b/tests/scripts/test_exporter.py index e25e4652e..bf81a6b22 100644 --- a/tests/scripts/test_exporter.py +++ b/tests/scripts/test_exporter.py @@ -76,6 +76,45 @@ def test_export_gaussian_splat_write_ply_mismatched_count(tmp_path: Path): ExportGaussianSplat.write_ply(filename, count, map_to_tensors) +def _read_ply_header(filename: str) -> str: + """Return the ASCII PLY header (up to and including the end_header line).""" + with open(filename, "rb") as f: + data = f.read() + end = data.index(b"end_header") + return data[:end].decode("ascii") + + +def _make_valid_tensors(count: int = 10) -> "OrderedDict[str, np.ndarray]": + return OrderedDict( + [ + ("x", np.random.rand(count).astype(np.float32)), + ("y", np.random.rand(count).astype(np.float32)), + ("z", np.random.rand(count).astype(np.float32)), + ] + ) + + +def test_export_gaussian_splat_write_ply_render_mode_comment(tmp_path: Path): + count = 10 + + # antialiased=True -> mip + mip_file = str(tmp_path / "mip.ply") + ExportGaussianSplat.write_ply(mip_file, count, _make_valid_tensors(count), antialiased=True) + assert "comment SplatRenderMode: mip" in _read_ply_header(mip_file) + + # antialiased=False -> default + default_file = str(tmp_path / "default.ply") + ExportGaussianSplat.write_ply(default_file, count, _make_valid_tensors(count), antialiased=False) + header = _read_ply_header(default_file) + assert "comment SplatRenderMode: default" in header + assert "comment SplatRenderMode: mip" not in header + + # antialiased=None (default) -> no SplatRenderMode comment (backwards compatible) + none_file = str(tmp_path / "none.ply") + ExportGaussianSplat.write_ply(none_file, count, _make_valid_tensors(count)) + assert "SplatRenderMode" not in _read_ply_header(none_file) + + if __name__ == "__main__": # Run the test test_export_gaussian_splat_write_ply(Path("."))