Align nerfacc renderers with nerfactory (#464)

add cuda support to renderers

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

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, :])