diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 2a2d30388..6570b8830 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -18,7 +18,6 @@ Implementation of Instant NGP. from __future__ import annotations -import math from dataclasses import dataclass, field from typing import Dict, List, Literal, Optional, Tuple, Type @@ -35,6 +34,7 @@ from nerfactory.configs import base as cfg from nerfactory.fields.instant_ngp_field import field_implementation_to_class from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.models.base import Model +from nerfactory.models.modules.ray_sampler import VolumetricSampler from nerfactory.optimizers.loss import MSELoss from nerfactory.utils import visualization from nerfactory.utils.callbacks import ( @@ -62,13 +62,15 @@ class InstantNGPModelConfig(cfg.ModelConfig): enable_density_field: bool = False """Whether to create a density field to filter samples.""" num_samples: int = 1024 - """Number of samples in field evaluation. Defaults to 1024,""" + """Number of samples in field evaluation.""" + grid_resolution: int = 128 + """Resolution of the grid used for the field.""" cone_angle: float = 0.0 """Should be set to 0.0 for blender scenes but 1./256 for real scenes.""" near_plane: float = 0.05 """How far along ray to start sampling.""" randomize_background: bool = False - """Whether to randomize the background color. Defaults to False.""" + """Whether to randomize the background color.""" class NGPModel(Model): @@ -95,29 +97,14 @@ class NGPModel(Model): # to match Ruilong's code naming self.scene_aabb = Parameter(self.scene_bounds.aabb.flatten(), requires_grad=False) - # setup some rendering settings - render_n_samples = 1024 - self.render_step_size = ( - (self.scene_aabb[3:] - self.scene_aabb[:3]).max() * math.sqrt(3) / render_n_samples - ).item() + # Sampler - # setup occupancy field with eval function - def occ_eval_fn(x: torch.Tensor) -> torch.Tensor: - """Evaluate occupancy given positions. - - Args: - x: positions with shape (N, 3). - Returns: - occupancy values with shape (N, 1). - """ - density_after_activation = self.field.density_fn(x) - # those two are similar when density is small. - # occupancy = 1.0 - torch.exp(-density_after_activation * render_step_size) - occupancy = density_after_activation * self.render_step_size - return occupancy - - # occupancy grid - self.occ_field = nerfacc.OccupancyField(occ_eval_fn=occ_eval_fn, aabb=self.scene_aabb, resolution=128) + self.sampler = VolumetricSampler( + aabb=self.scene_aabb, + density_fn=self.field.density_fn, + grid_resolution=self.config.grid_resolution, + num_samples=self.config.num_samples, + ) # losses self.rgb_loss = MSELoss() @@ -134,8 +121,8 @@ class NGPModel(Model): TrainingCallback( where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], update_every_num_iters=1, - func=self.occ_field.every_n_step, # will take in step - ) + func=self.sampler.occ_field.every_n_step, # will take in step + ), ] def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -149,6 +136,7 @@ class NGPModel(Model): assert self.field is not None def query_fn(positions: TensorType["bs", 3], directions: TensorType["bs", 3], only_density=False): + assert self.field is not None if only_density: return self.field.density_fn(positions) field_outputs = self.field.get_outputs_from_positions_and_direction(positions, directions) @@ -156,24 +144,65 @@ class NGPModel(Model): sigmas = field_outputs[FieldHeadNames.DENSITY] return rgbs, sigmas - ( - accumulated_color, - accumulated_depth, - accumulated_weight, - steps_counter, # pylint: disable=unused-variable - compact_steps_counter, # pylint: disable=unused-variable - ) = nerfacc.volumetric_rendering( - query_fn=query_fn, - rays_o=ray_bundle.origins, - rays_d=ray_bundle.directions, - scene_aabb=self.scene_aabb, - scene_occ_binary=self.occ_field.occ_grid_binary, - scene_resolution=self.occ_field.resolution, - render_bkgd=torch.ones(3, device=self.device), - render_step_size=self.render_step_size, - near_plane=self.config.near_plane, - stratified=self.training, # only use stratified sampling during training + n_rays = ray_bundle.origins.shape[0] + + with torch.no_grad(): + ray_samples, packed_info = self.sampler( + ray_bundle=ray_bundle, + near_plane=self.config.near_plane, + ) + + frustum_starts = ray_samples.frustums.starts + frustum_ends = ray_samples.frustums.ends + frustum_dirs = ray_samples.frustums.directions + frustum_positions = ray_samples.frustums.get_positions() + + # compat the samples thru volumetric rendering + with torch.no_grad(): + densities = self.field.density_fn(frustum_positions) + ( + compact_packed_info, + compact_frustum_starts, + compact_frustum_ends, + compact_frustum_positions, + compact_frustum_dirs, + ) = nerfacc.volumetric_rendering_steps( + packed_info, + densities, + frustum_starts, + frustum_ends, + frustum_positions, + frustum_dirs, + ) + + # network + compact_query_results = query_fn(compact_frustum_positions, compact_frustum_dirs) + compact_rgbs, compact_densities = compact_query_results[0], compact_query_results[1] + + # accumulation + compact_weights, compact_ray_indices = nerfacc.volumetric_rendering_weights( + compact_packed_info, + compact_densities, + compact_frustum_starts, + compact_frustum_ends, ) + accumulated_color = nerfacc.volumetric_rendering_accumulate( + compact_weights, compact_ray_indices, compact_rgbs, n_rays + ) + accumulated_weight = nerfacc.volumetric_rendering_accumulate(compact_weights, compact_ray_indices, None, n_rays) + accumulated_depth = nerfacc.volumetric_rendering_accumulate( + compact_weights, + compact_ray_indices, + (compact_frustum_starts + compact_frustum_ends) / 2.0, + n_rays, + ) + + if self.config.randomize_background: + # TODO(matt) Investigate randomizing per pixel not per batch. + render_bkgd = torch.rand(3, device=accumulated_color.device) + else: + render_bkgd = torch.ones(3, device=self.device) + accumulated_color = accumulated_color + render_bkgd * (1.0 - accumulated_weight) alive_ray_mask = accumulated_weight.squeeze(-1) > 0 diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 1105cf57f..a3c3569bf 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -16,6 +16,7 @@ Collection of sampling strategies """ +import math from abc import abstractmethod from typing import Callable, Optional, Tuple @@ -469,6 +470,8 @@ class VolumetricSampler(Sampler): super().__init__(num_samples=num_samples) self.aabb = aabb + self.render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item() + # setup occupancy field with eval function def occ_eval_fn(x: torch.Tensor) -> torch.Tensor: """Evaluate occupancy given positions. @@ -495,7 +498,6 @@ class VolumetricSampler(Sampler): def forward( self, ray_bundle: RayBundle, - render_step_size: float, num_samples: Optional[int] = None, near_plane: float = 0.0, ) -> Tuple[RaySamples, TensorType["total_samples", 3]]: @@ -503,7 +505,6 @@ class VolumetricSampler(Sampler): Args: ray_bundle: Rays to generate samples for - render_step_size: Step size for rendering. num_samples: Number of samples per ray near_plane: Near plane for raymarching @@ -516,7 +517,11 @@ class VolumetricSampler(Sampler): if self.density_field is not None: raise ValueError("NGPSpacedSampler does not support a density field") - num_samples = num_samples or self.num_samples + if num_samples is not None: + render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item() + else: + num_samples = self.num_samples + render_step_size = self.render_step_size rays_o = ray_bundle.origins.contiguous() rays_d = ray_bundle.directions.contiguous() @@ -532,7 +537,6 @@ class VolumetricSampler(Sampler): stratified=self.training, ) - # return samples zeros = torch.zeros_like(origins[:, :1]) ray_samples = RaySamples( frustums=Frustums(origins=origins, directions=dirs, starts=starts, ends=ends, pixel_area=zeros),