diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 52228cc4f..26503b714 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -35,7 +35,12 @@ 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.renderers.renderers import ( + AccumulationRenderer, + DepthRenderer, + RGBRenderer, +) +from nerfactory.utils import colors, visualization from nerfactory.utils.callbacks import ( TrainingCallback, TrainingCallbackAttributes, @@ -93,11 +98,9 @@ class NGPModel(Model): # torch or tiny-cuda-nn version self.field = field_implementation_to_class[self.config.field_implementation](self.scene_bounds.aabb) - # to match Ruilong's code naming self.scene_aabb = Parameter(self.scene_bounds.aabb.flatten(), requires_grad=False) # Sampler - self.sampler = VolumetricSampler( aabb=self.scene_aabb, density_fn=self.field.density_fn, @@ -105,6 +108,12 @@ class NGPModel(Model): num_samples=self.config.num_samples, ) + # renderers + background_color = None if self.config.randomize_background else colors.WHITE + self.renderer_rgb = RGBRenderer(background_color=background_color) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + # losses self.rgb_loss = MSELoss() @@ -143,38 +152,29 @@ class NGPModel(Model): ) 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( - 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, - (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2.0, - n_rays, + weights, ray_indices = nerfacc.volumetric_rendering_weights( + packed_info=packed_info, + sigmas=field_outputs[FieldHeadNames.DENSITY], + frustum_starts=ray_samples.frustums.starts, + frustum_ends=ray_samples.frustums.ends, ) - 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 + rgb = self.renderer_rgb( + rgb=field_outputs[FieldHeadNames.RGB], + weights=weights, + ray_indices=ray_indices, + num_rays=n_rays, + ) + depth = self.renderer_depth(weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=n_rays) + accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=n_rays) + alive_ray_mask = accumulation.squeeze(-1) > 0 outputs = { - "rgb": accumulated_color, - "accumulation": accumulated_weight, - "depth": accumulated_depth, + "rgb": rgb, + "accumulation": accumulation, + "depth": depth, "alive_ray_mask": alive_ray_mask, # the rays we kept from sampler } return outputs diff --git a/nerfactory/renderers/renderers.py b/nerfactory/renderers/renderers.py index 22cffd050..fa9d7a3b1 100644 --- a/nerfactory/renderers/renderers.py +++ b/nerfactory/renderers/renderers.py @@ -29,6 +29,7 @@ Example: import math from typing import Optional +import nerfacc import torch from torch import nn from torchtyping import TensorType @@ -54,6 +55,8 @@ class RGBRenderer(nn.Module): rgb: TensorType["bs":..., "num_samples", 3], weights: TensorType["bs":..., "num_samples", 1], background_color: Optional[TensorType[3]] = None, + ray_indices: Optional[TensorType["num_samples"]] = None, + num_rays: Optional[int] = None, ) -> TensorType["bs":..., 3]: """Composite samples along ray and render color image @@ -61,16 +64,23 @@ class RGBRenderer(nn.Module): rgb: RGB for each sample weights: Weights for each sample background_color: Background color as RGB. Defaults to random. + ray_indices: Ray index for each sample, used when samples are packed. + num_rays: Number of rays, used when samples are packed. Returns: Outputs rgb values. """ - rgb = torch.sum(weights * rgb, dim=-2) + if ray_indices is not None and num_rays is not None: + rgb = nerfacc.volumetric_rendering_accumulate(weights, ray_indices, rgb, num_rays) + accumulated_weight = nerfacc.volumetric_rendering_accumulate(weights, ray_indices, None, num_rays) + else: + rgb = torch.sum(weights * rgb, dim=-2) + accumulated_weight = torch.sum(weights, dim=-2) if background_color is None: background_color = torch.rand_like(rgb).to(rgb.device) - rgb = rgb + background_color.to(weights.device) * (1.0 - torch.sum(weights, dim=-2)) + rgb = rgb + background_color.to(weights.device) * (1.0 - accumulated_weight) return rgb @@ -78,18 +88,24 @@ class RGBRenderer(nn.Module): self, rgb: TensorType["bs":..., "num_samples", 3], weights: TensorType["bs":..., "num_samples", 1], + ray_indices: Optional[TensorType["num_samples"]] = None, + num_rays: Optional[int] = None, ) -> TensorType["bs":..., 3]: """Composite samples along ray and render color image Args: rgb: RGB for each sample weights: Weights for each sample + ray_indices: Ray index for each sample, used when samples are packed. + num_rays: Number of rays, used when samples are packed. Returns: Outputs of rgb values. """ - rgb = self.combine_rgb(rgb, weights, background_color=self.background_color) + rgb = self.combine_rgb( + rgb, weights, background_color=self.background_color, ray_indices=ray_indices, num_rays=num_rays + ) if not self.training: torch.clamp_(rgb, min=0.0, max=1.0) return rgb @@ -150,17 +166,24 @@ class AccumulationRenderer(nn.Module): def forward( cls, weights: TensorType["bs":..., "num_samples", 1], + ray_indices: Optional[TensorType["num_samples"]] = None, + num_rays: Optional[int] = None, ) -> TensorType["bs":..., 1]: """Composite samples along ray and calculate accumulation. Args: weights: Weights for each sample + ray_indices: Ray index for each sample, used when samples are packed. + num_rays: Number of rays, used when samples are packed. Returns: Outputs of accumulated values. """ - accumulation = torch.sum(weights, dim=-2) + if ray_indices is not None and num_rays is not None: + accumulation = nerfacc.volumetric_rendering_accumulate(weights, ray_indices, None, num_rays) + else: + accumulation = torch.sum(weights, dim=-2) return accumulation @@ -177,12 +200,20 @@ class DepthRenderer(nn.Module): raise ValueError(f"{method} is an invalid depth calculation method") self.method = method - def forward(self, weights: TensorType[..., "num_samples", 1], ray_samples: RaySamples) -> TensorType[..., 1]: + def forward( + self, + weights: TensorType[..., "num_samples", 1], + ray_samples: RaySamples, + ray_indices: Optional[TensorType["num_samples"]] = None, + num_rays: Optional[int] = None, + ) -> TensorType[..., 1]: """Composite samples along ray and calculate disparities. Args: weights: Weights for each sample. ray_samples: Set of ray samples. + ray_indices: Ray index for each sample, used when samples are packed. + num_rays: Number of rays, used when samples are packed. Returns: Outputs of depth values. @@ -191,7 +222,11 @@ class DepthRenderer(nn.Module): if self.method == "expected": eps = 1e-10 steps = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 - depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) + + if ray_indices is not None and num_rays is not None: + depth = nerfacc.volumetric_rendering_accumulate(weights, ray_indices, steps, num_rays) + else: + depth = torch.sum(weights * steps, dim=-2) / (torch.sum(weights, -2) + eps) depth = torch.clip(depth, steps[..., 0, :], steps[..., -1, :])