From 304a759e856f0949c30bb3a29de14a02c8e7daf0 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Wed, 29 Jun 2022 18:28:21 -0700 Subject: [PATCH] Tancik/scene contraction (#70) Contract the scene to a sphere of radius 2. This is the same contraction used in MipNeRF 360. Add mipnerf 360 graph and config. It currently doesn't contain the other advancements made in the paper (sample proposal network, regularizers) --- .vscode/launch.json | 10 + configs/graph_mipnerf_360.yaml | 30 +++ docs/conf.py | 1 + docs/reference/fields/modules/index.rst | 3 +- .../fields/modules/spatial_distortions.rst | 8 + environment/requirements.txt | 1 + pyrad/data/format/mipnerf_360.py | 10 +- pyrad/data/structs.py | 2 +- pyrad/fields/modules/spatial_distortions.py | 73 +++++++ pyrad/fields/nerf_field.py | 11 +- pyrad/graphs/mipnerf_360.py | 188 ++++++++++++++++++ 11 files changed, 325 insertions(+), 12 deletions(-) create mode 100644 configs/graph_mipnerf_360.yaml create mode 100644 docs/reference/fields/modules/spatial_distortions.rst create mode 100644 pyrad/fields/modules/spatial_distortions.py create mode 100644 pyrad/graphs/mipnerf_360.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 72a0b0534..93adebfe4 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -39,6 +39,16 @@ "--config-name=graph_instant_ngp.yaml" ] }, + { + "name": "Python: run_train mipnerf 360 (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "console": "integratedTerminal", + "args": [ + "--config-name=graph_mipnerf_360.yaml" + ] + }, { "name": "Python: run_train nerfw (gpu)", "type": "python", diff --git a/configs/graph_mipnerf_360.yaml b/configs/graph_mipnerf_360.yaml new file mode 100644 index 000000000..d5fdbaa7a --- /dev/null +++ b/configs/graph_mipnerf_360.yaml @@ -0,0 +1,30 @@ +defaults: + - graph_default + - _self_ + +experiment_name: mipnerf_360 +method_name: vanilla_nerf + +data: + # Train dataloading. + dataset_inputs_train: + data_directory: data/mipnerf_360/garden + dataset_format: mipnerf_360 + downscale_factor: 4 + # Eval dataloading + dataset_inputs_eval: + data_directory: data/mipnerf_360/garden + dataset_format: mipnerf_360 + downscale_factor: 8 + +trainer: + steps_per_test: 200 + +graph: + _target_: pyrad.graphs.mipnerf_360.MipNerf360Graph + num_coarse_samples: 128 + num_importance_samples: 128 + collider_config: + _target_: pyrad.graphs.modules.scene_colliders.NearFarCollider + near_plane: 0.5 + far_plane: 20.0 \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index cf795f9a8..6e1373447 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -43,6 +43,7 @@ extensions = [ "sphinx_copybutton", "nbsphinx", "sphinx_design", + "sphinx.ext.mathjax", ] # List of patterns, relative to source directory, that match files and diff --git a/docs/reference/fields/modules/index.rst b/docs/reference/fields/modules/index.rst index 3915476bf..096805081 100644 --- a/docs/reference/fields/modules/index.rst +++ b/docs/reference/fields/modules/index.rst @@ -11,4 +11,5 @@ TODO: High level description of field modules and how they connect together. embeddings field_heads field_losses - mlp \ No newline at end of file + mlp + spatial_distortions \ No newline at end of file diff --git a/docs/reference/fields/modules/spatial_distortions.rst b/docs/reference/fields/modules/spatial_distortions.rst new file mode 100644 index 000000000..b3be03895 --- /dev/null +++ b/docs/reference/fields/modules/spatial_distortions.rst @@ -0,0 +1,8 @@ +.. _spatial_distortions: + +Spatial Distortions +===================== + +.. automodule:: pyrad.fields.modules.spatial_distortions + :members: + :show-inheritance: \ No newline at end of file diff --git a/environment/requirements.txt b/environment/requirements.txt index 8065b6465..39bca324e 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -1,4 +1,5 @@ black[jupyter]==22.3.0 +functorch==0.1.1 h5py>=2.9.0 hydra-colorlog hydra-core diff --git a/pyrad/data/format/mipnerf_360.py b/pyrad/data/format/mipnerf_360.py index 1f3555b2a..75821776b 100644 --- a/pyrad/data/format/mipnerf_360.py +++ b/pyrad/data/format/mipnerf_360.py @@ -22,7 +22,7 @@ import imageio import numpy as np import torch -from pyrad.data.structs import DatasetInputs, SceneBounds +from pyrad.data.structs import DatasetInputs def load_mipnerf_360_data( @@ -31,7 +31,6 @@ def load_mipnerf_360_data( split: str = "train", val_skip: int = 8, auto_scale: bool = True, - aabb_scale: float = 5.0, ) -> DatasetInputs: """Processes mipnerf 360 data. Example data can be downloaded from https://jonbarron.info/mipnerf360/. @@ -109,18 +108,11 @@ def load_mipnerf_360_data( intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) intrinsics *= torch.tensor([cx, cy, focal_length]) - scene_bounds = SceneBounds( - aabb=torch.tensor( - [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 - ) - ) - dataset_inputs = DatasetInputs( image_filenames=image_filenames, downscale_factor=1, intrinsics=intrinsics, camera_to_world=camera_to_world, - scene_bounds=scene_bounds, ) return dataset_inputs diff --git a/pyrad/data/structs.py b/pyrad/data/structs.py index a87e608f9..9b6733d67 100644 --- a/pyrad/data/structs.py +++ b/pyrad/data/structs.py @@ -60,7 +60,7 @@ class SceneBounds: return length def get_center(self): - """Reterns the center of the box.""" + """Returns the center of the box.""" diff = self.aabb[1] - self.aabb[0] return self.aabb[0] + diff / 2.0 diff --git a/pyrad/fields/modules/spatial_distortions.py b/pyrad/fields/modules/spatial_distortions.py new file mode 100644 index 000000000..4b5566163 --- /dev/null +++ b/pyrad/fields/modules/spatial_distortions.py @@ -0,0 +1,73 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Space distortions.""" + +from typing import Union + +import torch +from functorch import jacrev, vmap +from torch import nn +from torchtyping import TensorType + +from pyrad.utils.math import Gaussians + + +class SpatialDistortion(nn.Module): + """Apply spatial distortions""" + + def forward(self, positions: Union[TensorType[..., 3], Gaussians]) -> Union[TensorType[..., 3], Gaussians]: + """ + Args: + positions (Union[TensorType[..., 3], Gaussians]): Sample to distort + + Returns: + Union[TensorType[..., 3], Gaussians]: distorted sample + """ + + +class SceneContraction(SpatialDistortion): + """Contract unbounded space into a sphere of radius 2. This contraction was proposed in MipNeRF-360. + We use the following contraction equation: + + .. math:: + + f(x) = \\begin{cases} + x & ||x|| \\leq 1 \\\\ + (2 - \\frac{1}{||x||})(\\frac{x}{||x||}) & ||x|| > 1 + \\end{cases} + """ + + def forward(self, positions: Union[TensorType[..., 3], Gaussians]) -> Union[TensorType[..., 3], Gaussians]: + def contract(x: TensorType[..., "in_dim"]) -> TensorType[..., "in_dim"]: + mag = x.norm(dim=-1) + mask = mag >= 1 + x[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None]) + + return x + + if isinstance(positions, Gaussians): + means = contract(positions.mean) + + contract = lambda x: (2 - (1 / x.norm(dim=-1, keepdim=True))) * (x / x.norm(dim=-1, keepdim=True)) + jc_means = vmap(jacrev(contract))(positions.mean) + + # Only update covariances on positions outside the unit sphere + mag = positions.mean.norm(dim=-1) + mask = mag >= 1 + positions.cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1) + + return Gaussians(mean=means, cov=positions.cov) + + return contract(positions) diff --git a/pyrad/fields/nerf_field.py b/pyrad/fields/nerf_field.py index d2b598c47..6a865fe84 100644 --- a/pyrad/fields/nerf_field.py +++ b/pyrad/fields/nerf_field.py @@ -26,6 +26,7 @@ from pyrad.fields.modules.field_heads import DensityFieldHead, FieldHead, FieldH from pyrad.fields.modules.mlp import MLP from pyrad.fields.base import Field from pyrad.cameras.rays import PointSamples +from pyrad.fields.modules.spatial_distortions import SpatialDistortion class NeRFField(Field): @@ -40,6 +41,7 @@ class NeRFField(Field): head_mlp_layer_width (int, optional): Width of output head MLP layers. Defaults to 128. skip_connections (Tuple, optional): Where to add skip connection in base MLP. Defaults to (4,). use_integrated_encoding (bool, optional): Used integrated samples as encoding input, Defaults to False. + spatial_distortion (SpatialDistortion, optional): Spatial distortion. Defaults to None. """ def __init__( @@ -53,11 +55,13 @@ class NeRFField(Field): skip_connections: Tuple[int] = (4,), field_heads: Tuple[FieldHead] = (RGBFieldHead(),), use_integrated_encoding: bool = False, + spatial_distortion: Optional[SpatialDistortion] = None, ) -> None: super().__init__() self.position_encoding = position_encoding self.direction_encoding = direction_encoding self.use_integrated_encoding = use_integrated_encoding + self.spatial_distortion = spatial_distortion self.mlp_base = MLP( in_dim=self.position_encoding.get_out_dim(), @@ -80,9 +84,14 @@ class NeRFField(Field): def get_density(self, point_samples: PointSamples): if self.use_integrated_encoding: gaussian_samples = point_samples.frustums.get_gaussian_blob() + if self.spatial_distortion is not None: + gaussian_samples = self.spatial_distortion(gaussian_samples) encoded_xyz = self.position_encoding(gaussian_samples.mean, covs=gaussian_samples.cov) else: - encoded_xyz = self.position_encoding(point_samples.frustums.get_positions()) + positions = point_samples.frustums.get_positions() + if self.spatial_distortion is not None: + positions = self.spatial_distortion(positions) + encoded_xyz = self.position_encoding(positions) base_mlp_out = self.mlp_base(encoded_xyz) density = self.field_output_density(base_mlp_out) return density, base_mlp_out diff --git a/pyrad/graphs/mipnerf_360.py b/pyrad/graphs/mipnerf_360.py new file mode 100644 index 000000000..e769d5998 --- /dev/null +++ b/pyrad/graphs/mipnerf_360.py @@ -0,0 +1,188 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Implementation of mip-NeRF. +""" + + +from typing import Dict, List + +import torch +from torch.nn import Parameter +from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity + +from pyrad.fields.modules.encoding import NeRFEncoding +from pyrad.fields.modules.field_heads import FieldHeadNames +from pyrad.fields.modules.spatial_distortions import SceneContraction +from pyrad.fields.nerf_field import NeRFField +from pyrad.graphs.base import Graph +from pyrad.optimizers.loss import MSELoss +from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler +from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer +from pyrad.utils import colors +from pyrad.cameras.rays import RayBundle +from pyrad.utils import visualization, writer + + +class MipNerf360Graph(Graph): + """mip-NeRF graph""" + + def __init__( + self, + intrinsics=None, + camera_to_world=None, + near_plane=2.0, + far_plane=6.0, + num_coarse_samples=64, + num_importance_samples=128, + **kwargs, + ) -> None: + self.near_plane = near_plane + self.far_plane = far_plane + self.num_coarse_samples = num_coarse_samples + self.num_importance_samples = num_importance_samples + self.field = None + super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + + def populate_fields(self): + """Set the fields.""" + + position_encoding = NeRFEncoding( + in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True + ) + direction_encoding = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) + + self.field = NeRFField( + position_encoding=position_encoding, + direction_encoding=direction_encoding, + use_integrated_encoding=True, + spatial_distortion=SceneContraction(), + ) + + 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, include_original=False) + + # renderers + self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) + self.renderer_accumulation = AccumulationRenderer() + self.renderer_depth = DepthRenderer() + + # losses + self.rgb_loss = MSELoss() + + # metrics + self.psnr = PeakSignalNoiseRatio(data_range=1.0) + self.ssim = StructuralSimilarityIndexMeasure() + self.lpips = LearnedPerceptualImagePatchSimilarity() + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + param_groups = {} + param_groups["fields"] = list(self.field.parameters()) + return param_groups + + def get_outputs(self, ray_bundle: RayBundle): + # uniform sampling + ray_samples_uniform = self.sampler_uniform(ray_bundle) + + # First pass: + field_outputs_coarse = self.field.forward(ray_samples_uniform.to_point_samples()) + weights_coarse = ray_samples_uniform.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY]) + rgb_coarse = self.renderer_rgb( + rgb=field_outputs_coarse[FieldHeadNames.RGB], + weights=weights_coarse, + ) + accumulation_coarse = self.renderer_accumulation(weights_coarse) + 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) + + # Second pass: + field_outputs_fine = self.field.forward(ray_samples_pdf.to_point_samples()) + weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY]) + rgb_fine = self.renderer_rgb( + rgb=field_outputs_fine[FieldHeadNames.RGB], + weights=weights_fine, + ) + accumulation_fine = self.renderer_accumulation(weights_fine) + depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf) + + outputs = { + "rgb_coarse": rgb_coarse, + "rgb_fine": rgb_fine, + "accumulation_coarse": accumulation_coarse, + "accumulation_fine": accumulation_fine, + "depth_coarse": depth_coarse, + "depth_fine": depth_fine, + } + return outputs + + def get_loss_dict(self, outputs, batch): + image = batch["image"] + rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) + rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) + loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} + return loss_dict + + def log_test_image_outputs(self, image_idx, step, batch, outputs): + image = batch["image"] + rgb_coarse = outputs["rgb_coarse"] + rgb_fine = outputs["rgb_fine"] + acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"]) + acc_fine = visualization.apply_colormap(outputs["accumulation_fine"]) + depth_coarse = visualization.apply_depth_colormap( + outputs["depth_coarse"], + accumulation=outputs["accumulation_coarse"], + near_plane=self.near_plane, + far_plane=self.far_plane, + ) + depth_fine = visualization.apply_depth_colormap( + outputs["depth_fine"], + accumulation=outputs["accumulation_fine"], + near_plane=self.near_plane, + far_plane=self.far_plane, + ) + + combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) + combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) + combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) + + writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) + writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) + writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) + + # 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, ...] + rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...] + + coarse_psnr = self.psnr(image, rgb_coarse) + fine_psnr = self.psnr(image, rgb_fine) + fine_ssim = self.ssim(image, rgb_fine) + fine_lpips = self.lpips(image, rgb_fine) + + writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step) + writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step) + writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step) + writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step) + + writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) + + return fine_psnr.item()