misc viewer fixes including when no samples in a chunk (#491)

viewer fixes

nerfactory/engine/trainer.py

@@ -126,6 +126,8 @@ class Trainer:
                 # if the viewer used, the rendering of the viewer will be included in the iteration train time
                 with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:
 
+                    self.pipeline.train()
+
                     # training callbacks before the training iteration
                     for callback in self.callbacks:
                         callback.run_callback_at_location(

nerfactory/fields/base.py

@@ -30,7 +30,7 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames
 class Field(nn.Module):
     """Base class for fields."""
 
-    def density_fn(self, positions: TensorType["bs":..., 3]):
+    def density_fn(self, positions: TensorType["bs":..., 3]) -> TensorType["bs":..., 1]:
         """Returns only the density. Used primarily with the density grid.
 
         Args:

nerfactory/models/modules/ray_sampler.py

@@ -441,6 +441,14 @@ class VolumetricSampler(Sampler):
             stratified=self.training,
         )
 
+        num_samples = starts.shape[0]
+        if num_samples == 0:
+            # create a single fake sample and update packed_info accordingly
+            # this says the last ray in packed_info has 1 sample, which starts and ends at 1
+            packed_info[-1, 1] = 1
+            starts = torch.ones((1, 1), dtype=starts.dtype, device=rays_o.device)
+            ends = torch.ones((1, 1), dtype=ends.dtype, device=rays_o.device)
+
         ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long()
         origins = rays_o[ray_indices]
         dirs = rays_d[ray_indices]

nerfactory/viewer/server/viewer_utils.py

@@ -487,7 +487,7 @@ class ViewerState:
                 vis_train_ratio = f"{int(vis_time / train_time * 100)}% spent on viewer"
                 self.vis["renderingState/vis_train_ratio"].write(vis_train_ratio)
             else:
-                self.vis["renderingState/vis_train_ratio"] = "Starting"
+                self.vis["renderingState/vis_train_ratio"].write("Starting")
         else:
             self.vis["renderingState/eval_fps"].write("Paused")
             self.vis["renderingState/train_eta"].write("Paused")