Faster instant-ngp (#166)

pyrad/graphs/instant_ngp.py

@@ -30,7 +30,7 @@ from pyrad.graphs.base import Graph
 from pyrad.optimizers.loss import MSELoss
 from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid
 from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
-from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler
+from pyrad.graphs.modules.ray_sampler import UniformSampler
 from pyrad.utils import colors
 from pyrad.cameras.rays import RayBundle
 from pyrad.utils import visualization, writer
@@ -67,10 +67,10 @@ class NGPGraph(Graph):
         self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb)
 
         # samplers
-        self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid)
-        # NOTE(ethan): are we sure we want the include_original flag used like this?
-        # it could be easily forgotten that it's by default True...?
-        self.sampler_pdf = PDFSampler(num_samples=128, include_original=False)
+        self.sampler_occupancy_grid = UniformSampler(num_samples=128)
+
+        # TODO stabalize occupancy grid.
+        # self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid)
 
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
@@ -93,30 +93,9 @@ class NGPGraph(Graph):
     def get_outputs(self, ray_bundle: RayBundle):
 
         # uniform sampling
-        ray_samples_uniform = self.sampler_occupancy_grid(ray_bundle)
-        field_outputs_uniform = self.field.forward(ray_samples_uniform)
-        weights_uniform = ray_samples_uniform.get_weights(field_outputs_uniform[FieldHeadNames.DENSITY])
-
-        # pdf sampling
-        ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_uniform)
-        field_outputs_pdf = self.field.forward(ray_samples_pdf)
-
-        # Hacky treatment of bins as points to allow us to merge uniform and pdf.
-        ts_uniform = (ray_samples_uniform.frustums.starts + ray_samples_uniform.frustums.ends) / 2.0
-        ts_pdf = (ray_samples_pdf.frustums.starts + ray_samples_pdf.frustums.ends) / 2.0
-        ts, indices = torch.sort(torch.cat([ts_uniform, ts_pdf], -2), -2)
-        bin_starts = ts
-        bin_ends = torch.cat([ts[..., 1:, :], ts[..., -1:, :]], dim=-2)
-        ray_samples = ray_bundle.get_ray_samples(bin_starts=bin_starts, bin_ends=bin_ends)
-
-        field_outputs = {}
-        for fo_name, _ in field_outputs_pdf.items():
-            fo_uniform = field_outputs_uniform[fo_name]
-            fo_pdf = field_outputs_pdf[fo_name]
-            fo_uniform_pdf = torch.cat([fo_uniform, fo_pdf], 1)
-            index = indices.view(fo_uniform_pdf[..., :1].shape)
-            index = index.expand(-1, -1, fo_uniform_pdf.shape[-1])  # TODO: don't hardcode this
-            field_outputs[fo_name] = torch.gather(fo_uniform_pdf, dim=1, index=index)
+        ray_samples = self.sampler_occupancy_grid(ray_bundle)
+        field_outputs = self.field.forward(ray_samples)
+        weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
 
         weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
         rgb = self.renderer_rgb(