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

@@ -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",

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

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

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
+   mlp
+   spatial_distortions

docs/reference/fields/modules/spatial_distortions.rst

@@ -0,0 +1,8 @@
+.. _spatial_distortions:
+
+Spatial Distortions
+=====================
+
+.. automodule:: pyrad.fields.modules.spatial_distortions
+   :members:
+   :show-inheritance:

environment/requirements.txt

@@ -1,4 +1,5 @@
 black[jupyter]==22.3.0
+functorch==0.1.1
 h5py>=2.9.0
 hydra-colorlog
 hydra-core

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

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
 

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)

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

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