Cleanup mipnerf implementation

.vscode/launch.json

@@ -36,7 +36,7 @@
             "program": "scripts/run_train.py",
             "console": "integratedTerminal",
             "args": [
-                "--config-name=instant_ngp.yaml"
+                "--config-name=graph_instant_ngp.yaml"
             ]
         },
         {

configs/graph_mipnerf.yaml

@@ -0,0 +1,11 @@
+defaults:
+  - graph_default
+  - _self_
+
+method_name: mipnerf
+
+graph:
+  network:
+    _target_: pyrad.graphs.mipnerf.MipNerfGraph
+    num_coarse_samples: 64
+    num_importance_samples: 128

pyrad/fields/modules/encoding.py

@@ -126,7 +126,7 @@ class NeRFEncoding(Encoding):
         if covs is None:
             encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1))
         else:
-            input_var = torch.diagonal(covs, dim1=-2, dim2=-1)[..., None, :] * freqs[:, None] ** 2
+            input_var = torch.diagonal(covs, dim1=-2, dim2=-1)[..., :, None] * freqs[None, :] ** 2
             input_var = input_var.reshape((*input_var.shape[:-2], -1))
             encoded_inputs = expected_sin(
                 torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1)
@@ -138,15 +138,16 @@ class NeRFEncoding(Encoding):
 
 
 class RFFEncoding(Encoding):
-    """Random Fourier Feature encoding"""
+    """Random Fourier Feature encoding. Supports integrated encodings."""
 
-    def __init__(self, in_dim: int, num_frequencies: int, scale: float) -> None:
+    def __init__(self, in_dim: int, num_frequencies: int, scale: float, include_input: bool = False) -> None:
         """
 
         Args:
             in_dim (int): Input dimension of tensor
             num_frequencies (int): Number of encoding frequencies
             scale (float): Std of Gaussian to sample frequencies. Must be greater than zero
+            include_input (float): Append the input coordinate to the encoding
         """
         super().__init__(in_dim)
 
@@ -156,20 +157,38 @@ class RFFEncoding(Encoding):
         self.scale = scale
         b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies))
         self.register_buffer(name="b_matrix", tensor=b_matrix)
+        self.include_input = include_input
 
     def get_out_dim(self) -> int:
         return self.num_frequencies * 2
 
-    def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
-        """
+    def forward(
+        self, in_tensor: TensorType[..., "input_dim"], covs: Optional[TensorType[..., "input_dim", "input_dim"]] = None
+    ) -> TensorType[..., "output_dim"]:
+        """Calculates RFF encoding. If covariances are provided the encodings will be integrated as proposed
+            in mip-NeRF.
+
         Args:
             in_tensor (TensorType[..., "input_dim"]): For best performance, the input tensor should be between 0 and 1.
+            covs (TensorType[..., "input_dim", "input_dim"], optional): Covariances of input points. Defaults to None.
 
         Returns:
             TensorType[..., "output_dim"]: Output values will be between -1 and 1
         """
+        in_tensor = 2 * torch.pi * in_tensor  # scale to [0, 2pi]
         scaled_inputs = in_tensor @ self.b_matrix  # [..., "num_frequencies"]
-        encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1)
+
+        if covs is None:
+            encoded_inputs = torch.sin(torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1))
+        else:
+            input_var = torch.sum(((covs @ self.b_matrix)) * self.b_matrix, -2)
+            encoded_inputs = expected_sin(
+                torch.cat([scaled_inputs, scaled_inputs + torch.pi / 2.0], dim=-1), torch.cat(2 * [input_var], dim=-1)
+            )
+
+        if self.include_input:
+            encoded_inputs = torch.cat([encoded_inputs, in_tensor], axis=-1)
+
         return encoded_inputs
 
 

pyrad/fields/nerf_field.py

@@ -41,6 +41,7 @@ class NeRFField(Field):
         head_mlp_layer_width: int = 128,
         skip_connections: Tuple[int] = (4,),
         field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
+        use_integrated_encoding: bool = False,
     ) -> None:
         """Create boilerplate NeRF field.
 
@@ -52,10 +53,12 @@ class NeRFField(Field):
             head_mlp_num_layers (int, optional): Number of layer for ourput head MLP. Defaults to 2.
             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.
         """
         super().__init__()
         self.position_encoding = position_encoding
         self.direction_encoding = direction_encoding
+        self.use_integrated_encoding = use_integrated_encoding
 
         self.mlp_base = MLP(
             in_dim=self.position_encoding.get_out_dim(),
@@ -77,7 +80,11 @@ class NeRFField(Field):
 
     def get_density(self, point_samples: PointSamples):
         """Computes and returns the densities."""
-        encoded_xyz = self.position_encoding(point_samples.frustums.get_positions())
+        if self.use_integrated_encoding:
+            gaussian_samples = point_samples.frustums.get_gaussian_blob()
+            encoded_xyz = self.position_encoding(gaussian_samples.mean, covs=gaussian_samples.cov)
+        else:
+            encoded_xyz = self.position_encoding(point_samples.frustums.get_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.py

@@ -0,0 +1,184 @@
+# 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.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 MipNerfGraph(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
+        )
+
+    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.bins)
+
+        # 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.bins)
+
+        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()

tests/fields/modules/test_encodings.py

@@ -48,6 +48,10 @@ def test_nerf_encoder():
     assert encoded.shape[-1] == out_dim
     assert torch.min(encoded) == 0
 
+    # Test integrated posenc
+    covs = torch.ones((2, 3, in_dim, in_dim))
+    encoded = encoder(in_tensor, covs=covs)
+
 
 def test_rff_encoder():
     """Test RFF encoder"""
@@ -64,6 +68,10 @@ def test_rff_encoder():
     encoded = encoder(in_tensor)
     assert encoded.shape[-1] == out_dim
 
+    # Test integrated encoding
+    covs = torch.ones((2, 3, in_dim, in_dim))
+    encoded = encoder(in_tensor, covs=covs)
+
 
 def test_tensor_vm_encoder():
     """Test TensorVM encoder"""