From cc79708df735f150a1c50c751a357cedee5ec2ae Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 9 Jun 2022 17:19:37 -0700 Subject: [PATCH] Cleanup mipnerf implementation --- .vscode/launch.json | 2 +- configs/graph_mipnerf.yaml | 11 ++ notebooks/visualize_encoders.ipynb | 148 ++++++++++++++++++-- pyrad/fields/modules/encoding.py | 31 ++++- pyrad/fields/nerf_field.py | 9 +- pyrad/graphs/mipnerf.py | 184 +++++++++++++++++++++++++ tests/fields/modules/test_encodings.py | 8 ++ 7 files changed, 377 insertions(+), 16 deletions(-) create mode 100644 configs/graph_mipnerf.yaml create mode 100644 pyrad/graphs/mipnerf.py diff --git a/.vscode/launch.json b/.vscode/launch.json index 983bee715..732a0399e 100644 --- a/.vscode/launch.json +++ b/.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" ] }, { diff --git a/configs/graph_mipnerf.yaml b/configs/graph_mipnerf.yaml new file mode 100644 index 000000000..e9b63e9e9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/notebooks/visualize_encoders.ipynb b/notebooks/visualize_encoders.ipynb index 644a999ca..129d91dff 100644 --- a/notebooks/visualize_encoders.ipynb +++ b/notebooks/visualize_encoders.ipynb @@ -28,7 +28,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 5, "metadata": {}, "outputs": [ { @@ -60,7 +60,65 @@ { "data": { "text/html": [ - "
" + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Encoded Integrate Values:\n", + "Covariance Magnitude: 0.01\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Covariance Magnitude: 0.1\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Covariance Magnitude: 1\n" + ] + }, + { + "data": { + "text/html": [ + "
" ], "text/plain": [ "" @@ -74,8 +132,9 @@ "num_frequencies = 4\n", "min_freq_exp = 0\n", "max_freq_exp = 6\n", - "include_input = True\n", + "include_input = False\n", "resolution = 128\n", + "covariance_magnitudes = [0.01, 0.1, 1]\n", "\n", "encoder = encoding.NeRFEncoding(\n", " in_dim=2,\n", @@ -93,7 +152,14 @@ "print(\"Input Values:\")\n", "media.show_images(torch.moveaxis(grid, 2, 0), cmap=\"plasma\", border=True)\n", "print(\"Encoded Values:\")\n", - "media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", border=True)" + "media.show_images(torch.moveaxis(encoded_values, 2, 0), vmin=-1, vmax=1, cmap=\"plasma\", border=True)\n", + "\n", + "print(\"Encoded Integrate Values:\")\n", + "for covariance_magnitude in covariance_magnitudes:\n", + " print(f\"Covariance Magnitude: {covariance_magnitude}\")\n", + " covs = torch.eye(2)[None, None, :, :] * covariance_magnitude\n", + " encoded_values = encoder(grid, covs=covs)\n", + " media.show_images(torch.moveaxis(encoded_values, 2, 0), vmin=-1, vmax=1, cmap=\"plasma\", border=True)" ] }, { @@ -106,7 +172,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 25, "metadata": {}, "outputs": [ { @@ -138,7 +204,65 @@ { "data": { "text/html": [ - "
" + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Encoded Integrate Values:\n", + "Covariance Magnitude: 0.001\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Covariance Magnitude: 0.01\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Covariance Magnitude: 0.1\n" + ] + }, + { + "data": { + "text/html": [ + "
" ], "text/plain": [ "" @@ -152,6 +276,7 @@ "num_frequencies = 8\n", "scale = 10\n", "resolution = 128\n", + "covariance_magnitudes = [0.001, 0.01, 0.1]\n", "\n", "encoder = encoding.RFFEncoding(in_dim=2, num_frequencies=num_frequencies, scale=scale)\n", "\n", @@ -163,7 +288,14 @@ "print(\"Input Values:\")\n", "media.show_images(torch.moveaxis(grid, 2, 0), cmap=\"plasma\", border=True)\n", "print(\"Encoded Values:\")\n", - "media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", border=True)" + "media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", vmin=-1, vmax=1, border=True)\n", + "\n", + "print(\"Encoded Integrate Values:\")\n", + "for covariance_magnitude in covariance_magnitudes:\n", + " print(f\"Covariance Magnitude: {covariance_magnitude}\")\n", + " covs = torch.eye(2)[None, None, :, :] * covariance_magnitude\n", + " encoded_values = encoder(grid, covs=covs)\n", + " media.show_images(torch.moveaxis(encoded_values, 2, 0), cmap=\"plasma\", vmin=-1, vmax=1, border=True)" ] }, { @@ -397,7 +529,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.13" + "version": "3.9.12" } }, "nbformat": 4, diff --git a/pyrad/fields/modules/encoding.py b/pyrad/fields/modules/encoding.py index 577c308ba..47724949e 100644 --- a/pyrad/fields/modules/encoding.py +++ b/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 diff --git a/pyrad/fields/nerf_field.py b/pyrad/fields/nerf_field.py index 6655670f8..94809625d 100644 --- a/pyrad/fields/nerf_field.py +++ b/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 diff --git a/pyrad/graphs/mipnerf.py b/pyrad/graphs/mipnerf.py new file mode 100644 index 000000000..2cbef0ee1 --- /dev/null +++ b/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() diff --git a/tests/fields/modules/test_encodings.py b/tests/fields/modules/test_encodings.py index 09c5caaa8..2053413c9 100644 --- a/tests/fields/modules/test_encodings.py +++ b/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"""