diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 1f7a0a8f9..df648d2df 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -137,19 +137,21 @@ class NeRFGraph(Graph): coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs # time end coarse network + coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) + coarse_renderer_outputs = self.renderer_rgb( rgb=coarse_field_outputs[FieldHeadNames.RGB], - density=coarse_field_outputs[FieldHeadNames.DENSITY], - deltas=uniform_ray_samples.deltas, + weights=coarse_weights, ) # RendererOutputs # fine network: - pdf_ray_samples = self.sampler_pdf(ray_bundle, uniform_ray_samples, coarse_field_outputs) # RaySamples + pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs + fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) + fine_renderer_outputs = self.renderer_rgb( rgb=fine_field_outputs[FieldHeadNames.RGB], - density=fine_field_outputs[FieldHeadNames.DENSITY], - deltas=pdf_ray_samples.deltas, + weights=fine_weights, ) # RendererOutputs # outputs: outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb} diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index dedf44506..887167d8b 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -2,14 +2,11 @@ Collection of renderers """ from dataclasses import dataclass -from typing import Set import torch from torch import nn from torchtyping import TensorType -from mattport.nerf.field_modules import field_heads - @dataclass class RendererOutputs: @@ -19,84 +16,25 @@ class RendererOutputs: density: TensorType["num_rays", 1] = None -class Renderer(nn.Module): - """Base Renderer. Intended to be subclassed""" - - def required_field_outputs(self) -> Set[field_heads.FieldHead]: - """ - Returns: - Set[RenderHead]: RenderHeads required for this renderer. - """ - raise NotImplementedError - - def forward( - self, - rgb: TensorType[..., "num_samples", 3], - density: TensorType[..., "num_samples", 1], - deltas: TensorType[..., "num_samples"], - ) -> RendererOutputs: - """Composite samples along ray and render image - - Args: - in_tensor (TensorType[..., "num_samples", -1]): feature for each sample - deltas (TensorType[..., "num_samples"]): depth of each sample - - Returns: - TensorType[..., "out_dim"]: Rendered ray - """ - raise NotImplementedError - - -class RGBRenderer(Renderer): +class RGBRenderer(nn.Module): """Standard volumetic rendering.""" - def required_field_outputs(self) -> Set[field_heads.FieldHead]: - return set(field_heads.RGBFieldHead, field_heads.DensityFieldHead) - def forward( self, rgb: TensorType[..., "num_samples", 3], - density: TensorType[..., "num_samples", 1], - deltas: TensorType[..., "num_samples"], + weights: TensorType[..., "num_samples"], ) -> RendererOutputs: """Composite samples along ray and render color image Args: - in_tensor (TensorType[..., "num_samples", -1]): RGB for each sample - deltas (TensorType[..., "num_samples"]): depth influence for each sample + rgb (TensorType[..., "num_samples", -1]): RGB for each sample + weights (TensorType[..., "num_samples"]): Weights for each sample Returns: TensorType[..., "out_dim"]: Composited RGB ray """ - delta_density = deltas * density[..., 0] - alphas = 1 - torch.exp(-delta_density) - - transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) - transmittance = torch.cat([torch.zeros((*transmittance.shape[:1], 1)).to(rgb.device), transmittance], axis=-1) - transmittance = torch.exp(-transmittance) # [..., "num_samples"] - - weights = alphas * transmittance # [..., "num_samples"] rgb = torch.sum(weights[..., None] * rgb, dim=-2) renderer_outputs = RendererOutputs(rgb=rgb) return renderer_outputs - - -class DepthRenderer(Renderer): - """_summary_ - - Args: - Renderer (_type_): _description_ - """ - - def required_field_outputs(self) -> Set[field_heads.FieldHead]: - return set(field_heads.DensityFieldHead) - - def forward( - self, - rgb: TensorType[..., "num_samples", 3], - density: TensorType[..., "num_samples", 1], - deltas: TensorType[..., "num_samples"], - ) -> RendererOutputs: - raise NotImplementedError diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 7915209da..c4b27af97 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -2,11 +2,10 @@ Collection of sampling strategies """ -from typing import Dict, Optional +from typing import Optional import torch from torch import nn from torchtyping import TensorType -from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.structures.rays import RayBundle, RaySamples @@ -89,9 +88,8 @@ class PDFSampler(nn.Module): def forward( self, - ray_bundle: RayBundle, coarse_ray_samples: RaySamples, - field_outputs: Dict[FieldHeadNames, TensorType], + weights: TensorType[..., "num_samples"], num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -108,27 +106,15 @@ class PDFSampler(nn.Module): Returns: RaySamples: Positions and deltas for samples along a ray """ + # TODO (matt) Look into torch.no_grad(), or torch.inference_mode num_samples = num_samples or self.num_samples - - # Calculate weight contributions along ray - # Todo(matt): This computation is duplicated - density = field_outputs[FieldHeadNames.DENSITY] - delta_density = coarse_ray_samples.deltas * density[..., 0] - alphas = 1 - torch.exp(-delta_density) - - transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) - transmittance = torch.cat( - [torch.zeros((*transmittance.shape[:1], 1)).to(density.device), transmittance], axis=-1 - ) - transmittance = torch.exp(-transmittance) # [..., "num_samples"] - weights = alphas * transmittance # [..., "num_samples"] - weights = weights[..., 1:] + weights = weights[..., :-1] # Add small offset to rays with zero weight to prevent NaNs weights_sum = torch.sum(weights, dim=-1, keepdim=True) padding = torch.relu(eps - weights_sum) - weights += padding / weights.shape[-1] + weights = weights + padding / weights.shape[-1] weights_sum += padding pdf = weights / weights_sum @@ -158,12 +144,15 @@ class PDFSampler(nn.Module): if self.include_original: ts, _ = torch.sort(torch.cat([coarse_ray_samples.ts, ts], -1), -1) + else: + ts, _ = torch.sort(ts, -1) + # Stop gradients ts = ts.detach() ray_samples = RaySamples( ts=ts, - ray_bundle=ray_bundle, + ray_bundle=coarse_ray_samples.ray_bundle, ) return ray_samples diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index b54af6488..6564ee876 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -92,3 +92,26 @@ class RaySamples: dists = torch.cat([dists, dists[..., -1:]], -1) # [N_rays, N_samples] deltas = dists * torch.norm(self.ray_bundle.directions[..., None, :], dim=-1) return deltas + + def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples"]: + """Return weights based on predicted densities + + Args: + densities (TensorType[..., "num_samples", 1]): Predicted densities for samples along ray + + Returns: + TensorType[..., "num_samples"]: Weights for each sample + """ + + delta_density = self.deltas * densities[..., 0] + alphas = 1 - torch.exp(-delta_density) + + transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) + transmittance = torch.cat( + [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 + ) + transmittance = torch.exp(-transmittance) # [..., "num_samples"] + + weights = alphas * transmittance # [..., "num_samples"] + + return weights diff --git a/tests/nerf/test_renderers.py b/tests/nerf/test_renderers.py index 7e2cd8c95..8c2170509 100644 --- a/tests/nerf/test_renderers.py +++ b/tests/nerf/test_renderers.py @@ -12,20 +12,15 @@ def test_rgb_renderer(): num_samples = 10 rgb = torch.ones((3, num_samples, 3)) - densities = torch.ones((3, num_samples, 1)) - deltas = torch.ones((3, num_samples)) + weights = torch.ones((3, num_samples)) + weights /= torch.sum(weights, axis=-1, keepdim=True) rgb_renderer = renderers.RGBRenderer() - assert rgb_renderer.required_field_outputs is not None - - out = rgb_renderer(rgb=rgb, density=densities, deltas=deltas) + out = rgb_renderer(rgb=rgb, weights=weights) assert torch.max(out.rgb) > 0.9 - out = rgb_renderer(rgb=rgb * 0, density=densities, deltas=deltas) - assert torch.max(out.rgb) == pytest.approx(0) - - out = rgb_renderer(rgb=rgb, density=densities * 0, deltas=deltas) + out = rgb_renderer(rgb=rgb * 0, weights=weights) assert torch.max(out.rgb) == pytest.approx(0) # TODO Tancik: Add more precise tests diff --git a/tests/nerf/test_sampler.py b/tests/nerf/test_sampler.py index 167bda16b..8e5ce7711 100644 --- a/tests/nerf/test_sampler.py +++ b/tests/nerf/test_sampler.py @@ -36,14 +36,14 @@ def test_pdf_sampler(): directions = torch.ones_like(origins) ray_bundle = RayBundle(origins=origins, directions=directions) - field_outputs = {FieldHeadNames.DENSITY: torch.ones((10, num_samples, 1))} - uniform_sampler = UniformSampler(near_plane=near_plane, far_plane=far_plane, num_samples=num_samples) coarse_ray_samples = uniform_sampler(ray_bundle) + weights = torch.ones((10, num_samples)) + # Just check that it doesn't crash. Need to add some actual tests. pdf_sampler = PDFSampler(num_samples) - pdf_sampler(ray_bundle, coarse_ray_samples, field_outputs, num_samples) + pdf_sampler(coarse_ray_samples, weights, num_samples) # TODO Tancik: Add more precise tests