add number of gaussians in logger; allow testing in nerfacto

nerfstudio/configs/base_config.py

@@ -88,6 +88,7 @@ class LocalWriterConfig(InstantiateConfig):
         writer.EventName.VIS_RAYS_PER_SEC,
         writer.EventName.TEST_RAYS_PER_SEC,
         writer.EventName.ETA,
+        writer.EventName.GAUSSIAN_NUM,
     )
     """specifies which stats will be logged/printed to terminal"""
     max_log_size: int = 10

nerfstudio/engine/trainer.py

@@ -36,6 +36,7 @@ from nerfstudio.configs.experiment_config import ExperimentConfig
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
 from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation
 from nerfstudio.engine.optimizers import Optimizers
+from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
 from nerfstudio.pipelines.base_pipeline import VanillaPipeline
 from nerfstudio.utils import profiler, writer
 from nerfstudio.utils.decorators import check_eval_enabled, check_main_thread, check_viewer_enabled
@@ -262,14 +263,21 @@ class Trainer:
 
                 # Skip the first two steps to avoid skewed timings that break the viewer rendering speed estimate.
                 if step > 1:
-                    writer.put_time(
-                        name=EventName.TRAIN_RAYS_PER_SEC,
-                        duration=self.world_size
-                        * self.pipeline.datamanager.get_train_rays_per_batch()
-                        / max(0.001, train_t.duration),
-                        step=step,
-                        avg_over_steps=True,
-                    )
+                    if isinstance(self.pipeline.model, GaussianSplattingModel):
+                        writer.put_scalar(
+                            name=EventName.GAUSSIAN_NUM,
+                            scalar=self.pipeline.model.get_gaussian_num,
+                            step=step,
+                        )
+                    else:
+                        writer.put_time(
+                            name=EventName.TRAIN_RAYS_PER_SEC,
+                            duration=self.world_size
+                            * self.pipeline.datamanager.get_train_rays_per_batch()
+                            / max(0.001, train_t.duration),
+                            step=step,
+                            avg_over_steps=True,
+                        )
 
                 self._update_viewer_state(step)
 

nerfstudio/models/gaussian_splatting.py

@@ -191,6 +191,10 @@ class GaussianSplattingModel(Model):
     def get_rotation(self):
         return self.rotation_activation(self._rotation)
     
+    @property
+    def get_gaussian_num(self):
+        return self._xyz.shape[0]
+    
     @property
     def get_xyz(self):
         return self._xyz

nerfstudio/utils/writer.py

@@ -56,6 +56,7 @@ class EventName(enum.Enum):
     TRAIN_RAYS_PER_SEC = "Train Rays / Sec"
     TEST_RAYS_PER_SEC = "Test Rays / Sec"
     VIS_RAYS_PER_SEC = "Vis Rays / Sec"
+    GAUSSIAN_NUM = "Number of Gaussians"
     CURR_TEST_PSNR = "Test PSNR"
 
 
@@ -468,7 +469,7 @@ class LocalWriter:
             if name in self.stats_to_track:
                 if "(time)" in name:
                     v = _format_time(v)
-                elif "Rays" in name:
+                elif "Rays" in name or "Gaussians" in name:
                     v = human_format(v)
                 else:
                     v = f"{v:0.4f}"

nerfstudio/viewer/server/render_state_machine.py

@@ -25,6 +25,7 @@ import torch
 from nerfstudio.cameras.cameras import Cameras
 from nerfstudio.model_components.renderers import \
     background_color_override_context
+from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
 from nerfstudio.utils import colormaps, writer
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
 from nerfstudio.viewer.server import viewer_utils
@@ -129,7 +130,9 @@ class RenderStateMachine(threading.Thread):
 
         with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext():
             #TODO jake-austin: Make this check whether the model inherits from a camera based model or a ray based model
-            if True:
+            #TODO Zhuoyang: First made some dummy judgements, need to be fixed later
+            isGaussianSplattingModel = isinstance(self.viewer.get_model(), GaussianSplattingModel)
+            if isGaussianSplattingModel:
                 pass
             else:
                 camera_ray_bundle = camera.generate_rays(camera_indices=0, aabb_box=self.viewer.get_model().render_aabb)
@@ -147,13 +150,13 @@ class RenderStateMachine(threading.Thread):
                             device=self.viewer.get_model().device,
                         )
                     with background_color_override_context(background_color), torch.no_grad():
-                        if True:
+                        if isGaussianSplattingModel:
                             outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera)
                         else:
                             outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)
                 else:
                     with torch.no_grad():
-                        if True:
+                        if isGaussianSplattingModel:
                             outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera)
                         else:
                             outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle)