Merge b7be44801f into 50e0e3c70c

nerfstudio/models/splatfacto.py

@@ -146,7 +146,9 @@ class SplatfactoModelConfig(ModelConfig):
     and apply them to the opacities of gaussians to preserve the total integrated density of splats.
 
     However, PLY exported with antialiased rasterize mode is not compatible with classic mode. Thus many web viewers that
-    were implemented for classic mode can not render antialiased mode PLY properly without modifications.
+    were implemented for classic mode can not render antialiased mode PLY properly without modifications. To help
+    compatible viewers pick the right mode, exported PLYs record this setting as a `SplatRenderMode` header comment
+    (`mip` for antialiased, `default` for classic).
     """
     camera_optimizer: CameraOptimizerConfig = field(default_factory=lambda: CameraOptimizerConfig(mode="off"))
     """Config of the camera optimizer to use"""
@@ -294,6 +296,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=self.config.rasterize_mode,
+            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:
@@ -539,7 +561,7 @@ class SplatfactoModel(Model):
 
         # apply the compensation of screen space blurring to gaussians
         if self.config.rasterize_mode not in ["antialiased", "classic"]:
-            raise ValueError("Unknown rasterize_mode: %s", self.config.rasterize_mode)
+            raise ValueError(f"Unknown rasterize_mode: {self.config.rasterize_mode}")
 
         if self.config.output_depth_during_training or not self.training:
             render_mode = "RGB+ED"

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[

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("."))