From 97cb3284687adc7e5d85afdc69111e675b6f6a20 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 18 May 2022 14:27:46 -0700 Subject: [PATCH] Refactor occupancy --- mattport/nerf/fields/density_field.py | 40 +++++++++++++++++ mattport/nerf/graph/base.py | 7 --- mattport/nerf/graph/vanilla_nerf.py | 21 +++------ mattport/nerf/occupancy_grid.py | 30 +++++++++++-- mattport/nerf/sampler.py | 65 +++++++++++++++++++++++++++ 5 files changed, 138 insertions(+), 25 deletions(-) create mode 100644 mattport/nerf/fields/density_field.py diff --git a/mattport/nerf/fields/density_field.py b/mattport/nerf/fields/density_field.py new file mode 100644 index 000000000..a31a0a67e --- /dev/null +++ b/mattport/nerf/fields/density_field.py @@ -0,0 +1,40 @@ +""" +Density Field. +""" + +from torch import nn + +from torchtyping import TensorType +from mattport.nerf.field_modules.field_heads import DensityFieldHead +from mattport.structures.rays import RaySamples + + +class DensityField(nn.Module): + """Base class for fields.""" + + def __init__(self, model: nn.Sequential, density_activation: nn.Module = nn.Softplus()) -> None: + super().__init__() + self.model = model + self.density_head = DensityFieldHead(in_dim=self.model[-1].get_out_dim(), activation=density_activation) + + def process_point(self, in_tensor: TensorType): + """_summary_""" + features = self.model(in_tensor) + density = self.density_head(features) + return density, features + + def proccess_ray(self, rays: RaySamples): + """_summary_""" + densities, features = self.process_point(in_tensor=rays.positions) + + weights = rays.get_weights(densities) + return densities, weights, features + + def forward(self, in_tensor): + """summary""" + return self.process_point(in_tensor) + + def density_fn(self, positions): + """Returns only the density.""" + density, _ = self.process_point(in_tensor=positions) + return density diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index f10c88213..af309e976 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -1,7 +1,6 @@ """ The Graph module contains all trainable parameters. """ -import logging from abc import abstractmethod from collections import defaultdict from typing import Dict, List, Union @@ -13,7 +12,6 @@ from torch.nn import Parameter from torchtyping import TensorType from mattport.nerf.dataset.structs import SceneBounds -from mattport.nerf.occupancy_grid import OccupancyGrid from mattport.nerf.ray_generator import RayGenerator from mattport.nerf.colliders import AABBBoxCollider from mattport.structures.cameras import get_camera_model @@ -65,7 +63,6 @@ class Graph(AbstractGraph): self.steps_per_occupancy_grid_update = steps_per_occupancy_grid_update self.kwargs = kwargs self.collider = None - self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb) self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world) self.populate_collider() self.populate_fields() @@ -97,10 +94,6 @@ class Graph(AbstractGraph): batch: Union[str, Dict[str, torch.tensor]] = None, step: int = None, ): - # update the occupancy grid - if step and step % self.steps_per_occupancy_grid_update == 0: - logging.info("Updating occupancy grid.") - self.occupancy_grid.update_occupancy_grid(density_fn=self.field_fine.density_fn) # get the rays original_ray_bundle = self.ray_generator.forward(ray_indices) # RayBundle diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index e548ccbe5..ffb90bea2 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -15,6 +15,7 @@ from mattport.nerf.field_modules.encoding import NeRFEncoding from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.fields.base import Field +from mattport.nerf.fields.density_field import DensityField from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer @@ -42,6 +43,9 @@ class NeRFField(Field): self.build_mlp_rgb() self.build_heads() + sequence = nn.Sequential(self.encoding_xyz, self.mlp_base) + self.density_field = DensityField(sequence) + def build_encodings(self): """Build the encodings.""" self.encoding_xyz = NeRFEncoding( @@ -89,9 +93,8 @@ class NeRFField(Field): ) if not valid_mask.any(): # empty mask return density, base_mlp_out - encoded_xyz = self.encoding_xyz(positions[valid_mask]) - base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz) - density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask]) + + density[valid_mask], base_mlp_out[valid_mask] = self.density_field(positions[valid_mask]) return density, base_mlp_out def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): @@ -138,7 +141,6 @@ class NeRFGraph(Graph): self.field_fine = NeRFField() def populate_misc_modules(self): - # samplers self.sampler_uniform = UniformSampler(num_samples=self.num_coarse_samples) self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples) @@ -162,7 +164,6 @@ class NeRFGraph(Graph): return param_groups def get_outputs(self, ray_bundle: RayBundle): - # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) @@ -189,12 +190,6 @@ class NeRFGraph(Graph): accumulation_fine = self.renderer_accumulation(weights_fine) depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf.ts) - # --- occupancy grid --- - # TODO(ethan): move this elsewhere. currently here for debugging purposes - density_occ = self.occupancy_grid.get_densities(ray_samples_pdf.positions) - weights_occ = ray_samples_pdf.get_weights(density_occ) - depth_occ = self.renderer_depth(weights_occ, ray_samples_pdf.ts) - outputs = { "rgb_coarse": rgb_coarse, "rgb_fine": rgb_fine, @@ -202,7 +197,6 @@ class NeRFGraph(Graph): "accumulation_fine": accumulation_fine, "depth_coarse": depth_coarse, "depth_fine": depth_fine, - "depth_occ": depth_occ, } return outputs @@ -241,9 +235,6 @@ class NeRFGraph(Graph): writer.write_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") writer.write_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") - depth_occ = visualization.apply_depth_colormap(outputs["depth_occ"]) - writer.write_image(name=f"image_idx_{image_idx}", image=depth_occ, step=step, group="depth_occupancy_grid") - # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...] diff --git a/mattport/nerf/occupancy_grid.py b/mattport/nerf/occupancy_grid.py index f44679c6a..4c4c7acc9 100644 --- a/mattport/nerf/occupancy_grid.py +++ b/mattport/nerf/occupancy_grid.py @@ -13,7 +13,9 @@ from mattport.utils.misc import is_not_none class OccupancyGrid(nn.Module): """Module to keep track of the density and occupancy.""" - def __init__(self, num_cascades=1, resolution=128, aabb=None) -> None: + def __init__( + self, num_cascades=1, resolution=128, aabb=None, density_fn=None, update_every_num_iters: int = 16 + ) -> None: super().__init__() assert is_not_none(aabb), "The axis-aligned bounding box aabb is not defined!" self.num_cascades = num_cascades # the number of levels (i.e, cascades) @@ -26,6 +28,10 @@ class OccupancyGrid(nn.Module): self.register_buffer("occupancy_grid", occupancy_grid) self.mean_density = 0.0 + self.density_fn = density_fn + self.update_every_num_iters = update_every_num_iters + self.iteration_count = 0 + def reset(self): """Zero out the occupancy grid.""" self.occupancy_grid.zero_() @@ -96,8 +102,18 @@ class OccupancyGrid(nn.Module): ) self.mean_density = torch.mean(self.occupancy_grid[valid_mask]).item() - def get_densities(self, xyzs: TensorType[..., 3]): - """Trilinear interpolation to get the density values.""" + @torch.no_grad() + def get_densities(self, xyzs: TensorType[..., 3], update_iter_count: bool = False) -> TensorType[..., 1]: + """Trilinear interpolation to get the density values. + + Args: + xyzs (TensorType[..., 3]): 3D querry coordinate + update_iter_count (bool, optional): Update current iteration. Used for knowing when to update grid. + Defaults to False. + + Returns: + TensorType[..., 1]: Density values + """ occupancy_grid = self.occupancy_grid[None, ...] # shape (1, num_cascades, X_res, Y_res, Z_res) xyzs_shape = xyzs.shape xyzs_reshaped = xyzs.view(1, -1, 1, 1, 3) @@ -112,6 +128,14 @@ class OccupancyGrid(nn.Module): padding_mode="zeros", ) densities = densities.view(*xyzs_shape[:-1], 1) + + # TODO we should be tracking the iteration another way + if self.training and update_iter_count: + self.iteration_count += 1 + if self.iteration_count % self.update_every_num_iters == 0: + if self.density_fn is not None: + self.update_occupancy_grid(self.density_fn) + return densities def forward(self, x): diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py index 70f1b94c0..0ddea9091 100644 --- a/mattport/nerf/sampler.py +++ b/mattport/nerf/sampler.py @@ -6,10 +6,75 @@ from typing import Optional import torch from torch import nn from torchtyping import TensorType +from mattport.nerf.occupancy_grid import OccupancyGrid from mattport.structures.rays import RayBundle, RaySamples +class OccupancyGridSampler(nn.Module): + """Sample based on occupancy grid""" + + def __init__( + self, + occupancy_grid: OccupancyGrid, + num_samples: int = 256, + weight_threshold: float = 1e-4, + train_stratified=True, + ) -> None: + """ + Args: + """ + super().__init__() + self.num_samples = num_samples + self.train_stratified = train_stratified + self.weight_threshold = weight_threshold + self.occupancy_grid = occupancy_grid + + @torch.no_grad() + def forward( + self, + ray_bundle: RayBundle, + num_samples: Optional[int] = None, + ) -> RaySamples: + """Generates position samples uniformly. + + Args: + ray_bundle (RayBundle): Rays to generate samples for + num_samples (Optional[int]): Number of samples per ray + + Returns: + RaySamples: Positions and deltas for samples along a ray + """ + assert ray_bundle.nears is not None + assert ray_bundle.fars is not None + + num_samples = num_samples or self.num_samples + num_rays = ray_bundle.origins.shape[0] + + bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device) # shape (num_samples+1,) + bins = ray_bundle.nears[:, None] + bins[None, :] * ( + ray_bundle.fars[:, None] - ray_bundle.nears[:, None] + ) # shape (num_rays, num_samples+1) + + if self.train_stratified and self.training: + t_rand = torch.rand((num_rays, num_samples), dtype=bins.dtype, device=bins.device) + ts = bins[:, :-1] + t_rand * (bins[:, 1:] - bins[:, :-1]) # shape (num_rays, num_samples) + else: + ts = (bins[:, 1:] + bins[:, :-1]) / 2 + + ray_samples = RaySamples( + ts=ts, + ray_bundle=ray_bundle, + ) + + densities = self.occupancy_grid.get_densities(ray_samples.positions, update_iter_count=True) + weights = ray_samples.get_weights(densities) + + valid_mask = weights >= self.weight_threshold + + return ray_samples, weights, valid_mask + + class UniformSampler(nn.Module): """Sample uniformly along a ray"""