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