Update bins

pyrad/cameras/rays.py

@@ -145,14 +145,16 @@ class RaySamples:
         frustums (Frustums): Frustums along ray.
         camera_indices (TensorType[..., 1]): Camera index.
         valid_mask (TensorType[...]): Rays that are valid.
-        bins (TensorType[..., 1]): frustum bins along ray.
+        bin_starts (TensorType[..., 1]): frustum starts along ray.
+        bin_ends (TensorType[..., 1]): frustum ends along ray.
         deltas )TensorType[..., 1]): "width" of each sample.
     """
 
     frustums: TensorType[..., 3]
     camera_indices: TensorType[..., 1] = None
     valid_mask: TensorType[...] = None
-    bins: TensorType[..., 1] = None
+    bin_starts: TensorType[..., 1] = None
+    bin_ends: TensorType[..., 1] = None
     deltas: TensorType[..., 1] = None
 
     def to_point_samples(self) -> PointSamples:
@@ -303,22 +305,25 @@ class RayBundle:
             camera_indices=camera_indices,
         )
 
-    def get_ray_samples(self, bins: TensorType["num_rays", "num_samples+1"]) -> RaySamples:
+    def get_ray_samples(
+        self, bin_starts: TensorType["num_rays", "num_samples"], bin_ends: TensorType["num_rays", "num_samples"]
+    ) -> RaySamples:
         """Produces samples for each ray by projection points along the ray direction.
 
 
         Args:
-            bins (TensorType["num_rays", "num_samples+1"]): Distance from origin of sample points.
+            bin_starts (TensorType["num_rays", "num_samples"]): Distance from origin to start of bin.
+            bin_ends (TensorType["num_rays", "num_samples"]): Distance from origin to end of bin.
 
         Returns:
             RaySamples: Samples projected along ray.
         """
         device = self.origins.device
-        num_samples = bins.shape[-1] - 1
+        num_samples = bin_starts.shape[-1]
 
-        valid_mask = torch.ones((bins.shape[0], num_samples), dtype=torch.bool, device=device)
+        valid_mask = torch.ones((bin_starts.shape), dtype=torch.bool, device=device)
 
-        dists = bins[:, 1:] - bins[:, :-1]  # [N_rays, N_samples]
+        dists = bin_ends - bin_starts  # [N_rays, N_samples]
         deltas = dists * torch.norm(self.directions[:, None, :], dim=-1)
 
         if is_not_none(self.camera_indices):
@@ -329,8 +334,8 @@ class RayBundle:
         frustums = Frustums(
             origins=self.origins.unsqueeze(1).repeat(1, num_samples, 1),
             directions=self.directions.unsqueeze(1).repeat(1, num_samples, 1),
-            frustum_starts=bins[:, :-1, None].to(device),
-            frustum_ends=bins[:, 1:, None].to(device),
+            frustum_starts=bin_starts[:, :, None].to(device),
+            frustum_ends=bin_ends[:, :, None].to(device),
             pixel_area=self.pixel_area.unsqueeze(1).repeat(1, num_samples, 1).to(device),
         )
 
@@ -338,7 +343,8 @@ class RayBundle:
             frustums=frustums,
             camera_indices=camera_indices,
             valid_mask=valid_mask,
-            bins=bins,
+            bin_starts=bin_starts,
+            bin_ends=bin_ends,
             deltas=deltas,
         )
 

pyrad/graphs/instant_ngp.py

@@ -89,8 +89,11 @@ class NGPGraph(Graph):
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_uniform)
         field_outputs_pdf = self.field.forward(ray_samples_pdf.to_point_samples())
 
-        ts, indices = torch.sort(torch.cat([ray_samples_uniform.bins, ray_samples_pdf.bins], -1), -1)
-        ray_samples = ray_bundle.get_ray_samples(ts)
+        # Hacky treatment of bins as points to allow us to merge uniform and pdf.
+        ts_uniform = (ray_samples_uniform.bin_starts + ray_samples_uniform.bin_ends) / 2.0
+        ts_pdf = (ray_samples_pdf.bin_starts + ray_samples_pdf.bin_ends) / 2.0
+        ts, indices = torch.sort(torch.cat([ts_uniform, ts_pdf], -1), -1)
+        ray_samples = ray_bundle.get_ray_samples(bin_starts=ts, bin_ends=ts)
         field_outputs = {}
         for fo_name, _ in field_outputs_pdf.items():
             fo_uniform = field_outputs_uniform[fo_name]
@@ -106,11 +109,11 @@ class NGPGraph(Graph):
             weights=weights,
         )
         accumulation = self.renderer_accumulation(weights)
-        depth = self.renderer_depth(weights, ray_samples.bins)
+        depth = self.renderer_depth(weights, ray_samples)
 
-        densities_occupancy_grid = self.occupancy_grid.get_densities(ray_samples.positions)
+        densities_occupancy_grid = self.occupancy_grid.get_densities(ray_samples.frustums.get_positions())
         weights_occupancy_grid = ray_samples.get_weights(densities_occupancy_grid)
-        depth_occupancy_grid = self.renderer_depth(weights_occupancy_grid, ray_samples.bins)
+        depth_occupancy_grid = self.renderer_depth(weights_occupancy_grid, ray_samples)
 
         outputs = {
             "rgb": rgb,

pyrad/graphs/mipnerf.py

@@ -105,7 +105,7 @@ class MipNerfGraph(Graph):
             weights=weights_coarse,
         )
         accumulation_coarse = self.renderer_accumulation(weights_coarse)
-        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.bins)
+        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)
 
         # pdf sampling
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
@@ -118,7 +118,7 @@ class MipNerfGraph(Graph):
             weights=weights_fine,
         )
         accumulation_fine = self.renderer_accumulation(weights_fine)
-        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.bins)
+        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
 
         outputs = {
             "rgb_coarse": rgb_coarse,

pyrad/graphs/modules/ray_sampler.py

@@ -109,7 +109,7 @@ class UniformSampler(Sampler):
             bin_lower = torch.cat([bins[..., :1], bin_centers], -1)
             bins = bin_lower + (bin_upper - bin_lower) * t_rand
 
-        ray_samples = ray_bundle.get_ray_samples(bins)
+        ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1], bin_ends=bins[..., 1:])
 
         return ray_samples
 
@@ -216,18 +216,28 @@ class PDFSampler(Sampler):
             x_end = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2)[0]
             return x_start, x_end
 
-        bins_g0, bins_g1 = find_interval(mask, ray_samples.bins)
+        # Force bins to not have a gap between them. Kinda hacky, should reconsider.
+        existing_bins = torch.cat(
+            [
+                ray_samples.bin_starts[..., :1],
+                (ray_samples.bin_starts[..., 1:] + ray_samples.bin_ends[..., :-1]) / 2.0,
+                ray_samples.bin_ends[..., -1:],
+            ],
+            axis=-1,
+        )
+
+        bins_g0, bins_g1 = find_interval(mask, existing_bins)
         cdf_g0, cdf_g1 = find_interval(mask, cdf)
 
         t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1)
         bins = bins_g0 + t * (bins_g1 - bins_g0)
 
         if self.include_original:
-            bins, _ = torch.sort(torch.cat([ray_samples.bins, bins], -1), -1)
+            bins, _ = torch.sort(torch.cat([existing_bins, bins], -1), -1)
 
         # Stop gradients
         bins = bins.detach()
 
-        ray_samples = ray_bundle.get_ray_samples(bins)
+        ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1], bin_ends=bins[..., 1:])
 
         return ray_samples

pyrad/graphs/nerfw.py

@@ -111,7 +111,7 @@ class NerfWGraph(Graph):
             rgb=field_outputs_coarse[FieldHeadNames.RGB],
             weights=weights_coarse,
         )
-        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.bins)
+        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)
 
         # pdf sampling
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
@@ -145,8 +145,8 @@ class NerfWGraph(Graph):
         density_transient = field_outputs_fine[FieldHeadNames.TRANSIENT_DENSITY]
 
         # depth
-        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.bins)
-        depth_fine_static = self.renderer_depth(weights_fine_static, ray_samples_pdf.bins)
+        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
+        depth_fine_static = self.renderer_depth(weights_fine_static, ray_samples_pdf)
 
         # uncertainty
         uncertainty = self.renderer_uncertainty(field_outputs_fine[FieldHeadNames.UNCERTAINTY], weights_fine_transient)

pyrad/graphs/semantic_nerf.py

@@ -135,7 +135,7 @@ class SemanticNerfGraph(NeRFGraph):
             weights=weights_coarse,
         )
         accumulation_coarse = self.renderer_accumulation(weights_coarse)
-        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.bins)
+        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)
 
         # pdf sampling
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
@@ -148,7 +148,7 @@ class SemanticNerfGraph(NeRFGraph):
             weights=weights_fine,
         )
         accumulation_fine = self.renderer_accumulation(weights_fine)
-        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.bins)
+        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
         semantic_fine = self.renderer_semantic(field_outputs_fine[FieldHeadNames.SEMANTICS], weights=weights_fine)
 
         outputs = {

pyrad/graphs/vanilla_nerf.py

@@ -114,7 +114,7 @@ class NeRFGraph(Graph):
             weights=weights_coarse,
         )
         accumulation_coarse = self.renderer_accumulation(weights_coarse)
-        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform.bins)
+        depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)
 
         # pdf sampling
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
@@ -127,7 +127,7 @@ class NeRFGraph(Graph):
             weights=weights_fine,
         )
         accumulation_fine = self.renderer_accumulation(weights_fine)
-        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.bins)
+        depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
 
         outputs = {
             "rgb_coarse": rgb_coarse,

pyrad/renderers/renderers.py

@@ -32,18 +32,19 @@ from typing import Optional
 import torch
 from torch import nn
 from torchtyping import TensorType
+from pyrad.cameras.rays import RaySamples
 
 from pyrad.utils.math import components_from_spherical_harmonics
 
 
 class RGBRenderer(nn.Module):
-    """Standard volumetic rendering."""
+    """Standard volumetic rendering.
+
+    Args:
+        background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
+    """
 
     def __init__(self, background_color: Optional[TensorType[3]] = None) -> None:
-        """
-        Args:
-            background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
-        """
         super().__init__()
         self.background_color = background_color
 
@@ -95,16 +96,16 @@ class RGBRenderer(nn.Module):
 
 
 class SHRenderer(nn.Module):
-    """Render RGB value from spherical harmonics."""
+    """Render RGB value from spherical harmonics.
+
+    Args:
+        background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
+        activation (Optional[nn.Module], optional): Output activation. Defaults to Sigmoid().
+    """
 
     def __init__(
         self, background_color: Optional[TensorType[3]] = None, activation: Optional[nn.Module] = nn.Sigmoid()
     ) -> None:
-        """
-        Args:
-            background_color (TensorType[3], optional): Background color as RGB. Defaults to random.
-            activation (Optional[nn.Module], optional): Output activation. Defaults to Sigmoid().
-        """
         super().__init__()
         self.background_color = background_color
         self.activation = activation
@@ -164,26 +165,24 @@ class AccumulationRenderer(nn.Module):
 
 
 class DepthRenderer(nn.Module):
-    """Calcualte depth along ray."""
+    """Calculate depth along ray.
+
+    Args:
+        method (str, optional): Depth calculation method. Defaults to 'expected'.
+    """
 
     def __init__(self, method: str = "expected") -> None:
-        """
-        Args:
-            method (str, optional): Depth calculation method. Defaults to 'expected'.
-        """
         super().__init__()
         if method not in {"expected"}:
             raise ValueError(f"{method} is an invalid depth calculation method")
         self.method = method
 
-    def forward(
-        self, weights: TensorType[..., "num_samples"], bins: TensorType[..., "num_samples+1"]
-    ) -> TensorType[..., 1]:
+    def forward(self, weights: TensorType[..., "num_samples"], ray_samples: RaySamples) -> TensorType[..., 1]:
         """Composite samples along ray and calculate disparities.
 
         Args:
-            weights (TensorType[..., "num_samples"]): Weights for each sample
-            bins (TensorType[..., "num_samples+1"]): Bins along ray
+            weights (TensorType[..., "num_samples"]): Weights for each sample.
+            ray_samples (RaySamples): Set of ray samples.
 
         Returns:
             TensorType[..., 1]: Outputs of depth values.
@@ -191,7 +190,7 @@ class DepthRenderer(nn.Module):
 
         if self.method == "expected":
             eps = 1e-10
-            steps = (bins[..., :-1] + bins[..., 1:]) / 2
+            steps = (ray_samples.bin_starts + ray_samples.bin_ends) / 2
             depth = torch.sum(weights * steps, dim=-1) / (torch.sum(weights, -1) + eps)
 
             depth = torch.clip(depth, steps[..., 0], steps[..., -1])

tests/renderers/test_renderers.py

@@ -3,6 +3,7 @@ Test renderers
 """
 import pytest
 import torch
+from pyrad.cameras.rays import Frustums, RaySamples
 
 from pyrad.renderers import renderers
 
@@ -62,12 +63,18 @@ def test_depth_renderer():
     weights = torch.ones((3, num_samples))
     weights /= torch.sum(weights, axis=-1, keepdim=True)
 
-    bins = torch.linspace(0, 100, num_samples + 1)
-    bins = torch.stack([bins, bins, bins], dim=0)
+    ray_samples = RaySamples(
+        frustums=Frustums.get_mock_frustum(),
+        camera_indices=torch.ones((num_samples)),
+        valid_mask=torch.ones((num_samples)),
+        bin_starts=torch.linspace(0, 100, num_samples),
+        bin_ends=torch.linspace(1, 101, num_samples),
+        deltas=torch.ones((num_samples)),
+    )
 
     depth_renderer = renderers.DepthRenderer()
 
-    depth = depth_renderer(weights=weights, bins=bins)
+    depth = depth_renderer(weights=weights, ray_samples=ray_samples)
     assert torch.min(depth) > 0