From 201f85a175f2781953317312340b8a8049c78dc1 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 10 Jun 2022 16:45:27 -0700 Subject: [PATCH] Update bins --- pyrad/cameras/rays.py | 26 ++++++++++------- pyrad/graphs/instant_ngp.py | 13 +++++---- pyrad/graphs/mipnerf.py | 4 +-- pyrad/graphs/modules/ray_sampler.py | 18 +++++++++--- pyrad/graphs/nerfw.py | 6 ++-- pyrad/graphs/semantic_nerf.py | 4 +-- pyrad/graphs/vanilla_nerf.py | 4 +-- pyrad/renderers/renderers.py | 43 ++++++++++++++--------------- tests/renderers/test_renderers.py | 13 +++++++-- 9 files changed, 78 insertions(+), 53 deletions(-) diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index a88bb7f82..d2ac1996d 100644 --- a/pyrad/cameras/rays.py +++ b/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, ) diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index b20e67537..8708d1969 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/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, diff --git a/pyrad/graphs/mipnerf.py b/pyrad/graphs/mipnerf.py index 2cbef0ee1..6dd68e5b8 100644 --- a/pyrad/graphs/mipnerf.py +++ b/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, diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index 0c74af005..f8c035be6 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/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 diff --git a/pyrad/graphs/nerfw.py b/pyrad/graphs/nerfw.py index 0c8f054bf..1f4c234fc 100644 --- a/pyrad/graphs/nerfw.py +++ b/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) diff --git a/pyrad/graphs/semantic_nerf.py b/pyrad/graphs/semantic_nerf.py index ccbec4ace..e3121e5ac 100644 --- a/pyrad/graphs/semantic_nerf.py +++ b/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 = { diff --git a/pyrad/graphs/vanilla_nerf.py b/pyrad/graphs/vanilla_nerf.py index 5819332b9..e5beeff1f 100644 --- a/pyrad/graphs/vanilla_nerf.py +++ b/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, diff --git a/pyrad/renderers/renderers.py b/pyrad/renderers/renderers.py index a966a6db2..e4ece1677 100644 --- a/pyrad/renderers/renderers.py +++ b/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]) diff --git a/tests/renderers/test_renderers.py b/tests/renderers/test_renderers.py index 67f780b57..85d7121f2 100644 --- a/tests/renderers/test_renderers.py +++ b/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