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()