From fe66547f2429df9c5f64e0b1ae46b0fc8abe3c2c Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Sat, 17 Sep 2022 18:17:58 -0700 Subject: [PATCH] Cleanup nerfacc integration (#463) * cleanup nerfacc implementation * typing --- nerfactory/fields/instant_ngp_field.py | 20 --------- nerfactory/models/instant_ngp.py | 53 ++++-------------------- nerfactory/models/modules/ray_sampler.py | 24 ++++++++++- 3 files changed, 32 insertions(+), 65 deletions(-) diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 2c57a5c80..779563249 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -151,26 +151,6 @@ class TCNNInstantNGPField(Field): rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) return {FieldHeadNames.RGB: rgb} - def get_outputs_from_positions_and_direction(self, positions: TensorType["bs", 3], directions: TensorType["bs", 3]): - """Given positions and directions, return the outputs of the field. - - Args: - positions: positions of the points - directions: directions at the point locations - """ - # assume flat positions and directions - # TODO: normalize positions - positions = SceneBounds.get_normalized_positions(positions, self.aabb) - directions = get_normalized_directions(directions) - h = self.mlp_base(positions) - density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) - density = trunc_exp(density_before_activation.to(positions)) - # TODO: normalize directions - d = self.direction_encoding(directions) - h = torch.cat([d, base_mlp_out], dim=-1) - rgb = self.mlp_head(h).to(directions) - return {FieldHeadNames.RGB: rgb, FieldHeadNames.DENSITY: density} - class TorchInstantNGPField(NeRFField): """ diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 6570b8830..52228cc4f 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -27,7 +27,6 @@ from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity -from torchtyping import TensorType from nerfactory.cameras.rays import RayBundle from nerfactory.configs import base as cfg @@ -135,15 +134,6 @@ class NGPModel(Model): def get_outputs(self, ray_bundle: RayBundle): 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) - rgbs = field_outputs[FieldHeadNames.RGB] - sigmas = field_outputs[FieldHeadNames.DENSITY] - return rgbs, sigmas - n_rays = ray_bundle.origins.shape[0] with torch.no_grad(): @@ -152,48 +142,23 @@ class NGPModel(Model): 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] + field_outputs = self.field(ray_samples) + rgbs = field_outputs[FieldHeadNames.RGB] + densities = field_outputs[FieldHeadNames.DENSITY] # 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 + packed_info, + densities, + ray_samples.frustums.starts, + ray_samples.frustums.ends, ) + accumulated_color = nerfacc.volumetric_rendering_accumulate(compact_weights, compact_ray_indices, 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, + (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2.0, n_rays, ) diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index a3c3569bf..51a80e914 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -469,6 +469,7 @@ class VolumetricSampler(Sampler): """ super().__init__(num_samples=num_samples) self.aabb = aabb + self.density_fn = density_fn self.render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item() @@ -500,6 +501,7 @@ class VolumetricSampler(Sampler): ray_bundle: RayBundle, num_samples: Optional[int] = None, near_plane: float = 0.0, + remove_occluded_samples: bool = True, ) -> Tuple[RaySamples, TensorType["total_samples", 3]]: """Generate ray samples in a bounding box. @@ -507,6 +509,7 @@ class VolumetricSampler(Sampler): ray_bundle: Rays to generate samples for num_samples: Number of samples per ray near_plane: Near plane for raymarching + remove_occluded_samples: Whether to remove occluded samples. (requires evaluating the density function) Returns: a tuple of (ray_samples, packed_info) @@ -537,8 +540,27 @@ class VolumetricSampler(Sampler): stratified=self.training, ) + if remove_occluded_samples: + with torch.no_grad(): + positions = origins + dirs * (starts + ends) / 2.0 + densities = self.density_fn(positions) + packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps( + packed_info, + densities, + starts, + ends, + origins, + dirs, + ) + zeros = torch.zeros_like(origins[:, :1]) ray_samples = RaySamples( - frustums=Frustums(origins=origins, directions=dirs, starts=starts, ends=ends, pixel_area=zeros), + frustums=Frustums( + origins=origins, + directions=dirs, + starts=starts, + ends=ends, + pixel_area=zeros, + ), ) return ray_samples, packed_info