diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 83b5de593..8f3c5f20c 100644 --- a/nerfactory/engine/trainer.py +++ b/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( diff --git a/nerfactory/fields/base.py b/nerfactory/fields/base.py index af8778d76..a1bba824d 100644 --- a/nerfactory/fields/base.py +++ b/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: diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 9116d9b8e..9411850e3 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/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] diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index 6893ad68a..fcc69770e 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/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")