From cd6342733ba381bee6369ea218f252d735f6d949 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 16 May 2022 12:27:33 -0700 Subject: [PATCH] Add spherical harmoincs --- mattport/nerf/field_modules/encoding.py | 24 +++++ mattport/nerf/field_modules/field_heads.py | 16 +++ mattport/nerf/renderers.py | 61 +++++++++++ mattport/utils/math.py | 69 ++++++++++++ notebooks/visualize_encoders.ipynb | 116 ++++++++++++++++++++- tests/nerf/field_modules/test_encodings.py | 20 ++++ tests/nerf/test_renderers.py | 50 ++++++++- 7 files changed, 353 insertions(+), 3 deletions(-) create mode 100644 mattport/utils/math.py diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index aa29a7945..2eb0b7375 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -9,6 +9,7 @@ from torch import nn from torchtyping import TensorType from mattport.nerf.field_modules.base import FieldModule +from mattport.utils.math import components_from_spherical_harmonics class Encoding(FieldModule): @@ -348,3 +349,26 @@ class TensorVMEncoding(Encoding): ) self.resolution = resolution + + +class SHEncoding(Encoding): + """Spherical harmonic encoding""" + + def __init__(self, levels: int = 4) -> None: + """ + Args: + levels (int, optional): Number of spherical hamonic levels to encode. Defaults to 4. + """ + super().__init__(in_dim=3) + + if levels <= 0 or levels > 4: + raise ValueError(f"Spherical harmonic encoding only suports 1 to 4 levels, requested {levels}") + + self.levels = levels + + def get_out_dim(self) -> int: + return self.levels**2 + + @torch.no_grad() + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor) diff --git a/mattport/nerf/field_modules/field_heads.py b/mattport/nerf/field_modules/field_heads.py index 5af4ded2a..7f49ab341 100644 --- a/mattport/nerf/field_modules/field_heads.py +++ b/mattport/nerf/field_modules/field_heads.py @@ -75,3 +75,19 @@ class RGBFieldHead(FieldHead): def __init__(self, in_dim: int, activation: Optional[nn.Module] = nn.Sigmoid()) -> None: super().__init__(in_dim, out_dim=3, field_head_name=FieldHeadNames.RGB, activation=activation) + + +class SHFieldHead(FieldHead): + """Spherical harmonics output""" + + def __init__(self, in_dim: int, levels: int = 3, channels: int = 3, activation: Optional[nn.Module] = None) -> None: + """_summary_ + + Args: + in_dim (int): Input dimension + levels (int, optional): Number of spherical harmonics layers. Defaults to 3. + channels (int, optional): Number of channels. Defaults to 3 (ie RGB). + activation (Optional[nn.Module], optional): Output activation. Defaults to None. + """ + out_dim = channels * levels**2 + super().__init__(in_dim, out_dim=out_dim, field_head_name=FieldHeadNames.RGB, activation=activation) diff --git a/mattport/nerf/renderers.py b/mattport/nerf/renderers.py index f746b5942..740beaea2 100644 --- a/mattport/nerf/renderers.py +++ b/mattport/nerf/renderers.py @@ -2,12 +2,15 @@ Collection of renderers """ from dataclasses import dataclass +import math from typing import Optional import torch from torch import nn from torchtyping import TensorType +from mattport.utils.math import components_from_spherical_harmonics + @dataclass class RendererOutputs: @@ -60,6 +63,64 @@ class RGBRenderer(nn.Module): return renderer_outputs +class SHRenderer(nn.Module): + """Render RGB value from spherical harmonics.""" + + def __init__( + self, activation: Optional[nn.Module] = nn.Sigmoid(), background_color: Optional[TensorType[3]] = None + ) -> None: + """ + Args: + background_color (TensorType[3], optional): Background color as RGB. Defaults to black. + activation (Optional[nn.Module], optional): Output activation. Defaults to Sigmoid(). + """ + super().__init__() + self.background_color = background_color + self.activation = activation + + def forward( + self, + sh: TensorType[..., "num_samples", "coeffs"], + directions: TensorType[..., "num_samples", 3], + weights: TensorType[..., "num_samples"], + ) -> RendererOutputs: + """Composite samples along ray and render color image + + Args: + sh (TensorType[..., "num_samples", "coeffs"]): Spherical hamonics coefficients for each sample + directions: (TensorType[..., "num_samples", 3]): Sample direction + weights (TensorType[..., "num_samples"]): Weights for each sample + + Returns: + RendererOutputs: Outputs with rgb values. + """ + + sh = sh.view(*sh.shape[:-1], 3, sh.shape[-1] // 3) + + levels = int(math.sqrt(sh.shape[-1])) + components = components_from_spherical_harmonics(levels=levels, directions=directions) + + rgb = sh * components[..., None, :] # [..., num_samples, 3, sh_components] + rgb = torch.sum(sh, dim=-1) # [..., num_samples, 3] + + rgb = torch.sum(weights[..., None] * rgb, dim=-2) + 0.5 + + if self.activation is not None: + self.activation(rgb) + + if self.background_color is not None: + rgb = rgb + self.background_color.to(weights.device)[None, ...] * ( + 1.0 - torch.sum(weights, dim=-1)[..., None] + ) + torch.clamp_(rgb, min=0.0, max=1.0) + + assert torch.max(rgb) <= 1.0 + assert torch.min(rgb) >= 0.0 + + renderer_outputs = RendererOutputs(rgb=rgb) + return renderer_outputs + + class AccumulationRenderer(nn.Module): """Accumulated value along a ray.""" diff --git a/mattport/utils/math.py b/mattport/utils/math.py new file mode 100644 index 000000000..ffd66e3c0 --- /dev/null +++ b/mattport/utils/math.py @@ -0,0 +1,69 @@ +""" Math Helper Functions """ + +import torch +from torchtyping import TensorType + + +def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]: + """ + Args: + levels (int): Number of spherical harmonic levels to compute. + coeffs (TensorType[..., "coeffs"]): Spherical hamonic coefficients + + Returns: + TensorType[..., "components"]: Value for each component. + """ + num_components = levels**2 + components = torch.zeros((*directions.shape[:-1], num_components)).to(directions.device) + + assert levels > 0 and levels <= 5, f"SH levels must be in [1,4], got {levels}" + assert directions.shape[-1] == 3, f"Direction input should have three dimensions. Got {directions.shape[-1]}" + + x = directions[..., 0] + y = directions[..., 1] + z = directions[..., 2] + + xx = x**2 + yy = y**2 + zz = z**2 + + # l0 + components[..., 0] = 0.28209479177387814 + + # l1 + if levels > 1: + components[..., 1] = 0.4886025119029199 * y + components[..., 2] = 0.4886025119029199 * z + components[..., 3] = 0.4886025119029199 * x + + # l2 + if levels > 2: + components[..., 4] = 1.0925484305920792 * x * y + components[..., 5] = 1.0925484305920792 * y * z + components[..., 6] = 0.9461746957575601 * zz + components[..., 7] = 1.0925484305920792 * x * z + components[..., 8] = 0.5462742152960396 * (xx - yy) + + # l3 + if levels > 3: + components[..., 9] = 0.5900435899266435 * y * (3 * xx - yy) + components[..., 10] = 2.890611442640554 * x * y * z + components[..., 11] = 0.4570457994644658 * y * (5 * zz - 1) + components[..., 12] = 0.3731763325901154 * z * (5 * zz - 3) + components[..., 13] = 0.4570457994644658 * x * (5 * zz - 1) + components[..., 14] = 1.445305721320277 * z * (xx - yy) + components[..., 15] = 0.5900435899266435 * x * (xx - 3 * yy) + + # l4 + if levels > 4: + components[..., 16] = 2.5033429417967046 * x * y * (xx - yy) + components[..., 17] = 1.7701307697799304 * y * z * (3 * xx - yy) + components[..., 18] = 0.9461746957575601 * x * y * (7 * zz - 1) + components[..., 19] = 0.6690465435572892 * y * (7 * zz - 3 * z) + components[..., 20] = 0.10578554691520431 * (35 * zz * zz - 30 * zz + 3) + components[..., 21] = 0.6690465435572892 * x * z * (7 * zz - 3) + components[..., 22] = 0.47308734787878004 * (xx - yy) * (7 * zz - 1) + components[..., 23] = 1.7701307697799304 * x * z * (xx - 3 * yy) + components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) + + return components diff --git a/notebooks/visualize_encoders.ipynb b/notebooks/visualize_encoders.ipynb index 2fde5f60c..618d23412 100644 --- a/notebooks/visualize_encoders.ipynb +++ b/notebooks/visualize_encoders.ipynb @@ -8,7 +8,7 @@ "source": [ "import torch\n", "import mediapy as media\n", - "from mattport.nerf.field_modules import encoding\n" + "from mattport.nerf.field_modules import encoding" ] }, { @@ -251,6 +251,118 @@ "media.show_images(encoded_images.detach().numpy(), cmap=\"plasma\", border=True)\n" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Spherical Harmonic Encoding\n", + "Encode direction using spherical harmoincs. (Mostly used to encode viewing direction)" + ] + }, + { + "cell_type": "code", + "execution_count": 66, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Level: 1\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Level: 2\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Level: 3\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Level: 4\n" + ] + }, + { + "data": { + "text/html": [ + "
" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "levels = 4\n", + "\n", + "height = 100\n", + "width = 150\n", + "\n", + "encoder = encoding.SHEncoding(levels=levels)\n", + "\n", + "theta = torch.linspace(-torch.pi, torch.pi, width)\n", + "phi = torch.linspace(0, torch.pi, height)\n", + "[theta, phi] = torch.meshgrid([theta, phi], indexing=\"xy\")\n", + "\n", + "directions = torch.stack([torch.cos(theta) * torch.sin(phi), torch.sin(theta) * torch.sin(phi), torch.cos(phi)], dim=-1)\n", + "\n", + "encoded_values = encoder(directions)\n", + "encoded_values = torch.moveaxis(encoded_values, 2, 0)\n", + "\n", + "for level in range(levels):\n", + " print(f\"Level: {level+1}\")\n", + " media.show_images(encoded_values[level**2:(level+1)**2,...], cmap=\"plasma\", border=True)" + ] + }, { "cell_type": "code", "execution_count": null, @@ -278,7 +390,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.11" + "version": "3.9.12" }, "orig_nbformat": 4 }, diff --git a/tests/nerf/field_modules/test_encodings.py b/tests/nerf/field_modules/test_encodings.py index 572161f7f..1708330d8 100644 --- a/tests/nerf/field_modules/test_encodings.py +++ b/tests/nerf/field_modules/test_encodings.py @@ -1,6 +1,7 @@ """ Encoding Tests """ +from multiprocessing.sharedctypes import Value import pytest import torch from mattport.nerf.field_modules import encoding @@ -111,9 +112,28 @@ def test_tensor_cp_encoder(): encoder.upsample_grid(resolution=64) +def test_tensor_sh_encoder(): + """Test Spherical Harmonic encoder""" + + levels = 4 + out_dim = levels**2 + + with pytest.raises(ValueError): + encoder = encoding.SHEncoding(levels=5) + + encoder = encoding.SHEncoding(levels=levels) + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.zeros((10, 3)) + in_tensor[..., 1] = 1 + encoded = encoder.encode(in_tensor) + assert encoded.shape == (10, out_dim) + + if __name__ == "__main__": test_scaling_and_offset() test_nerf_encoder() test_rff_encoder() test_tensor_vm_encoder() test_tensor_cp_encoder() + test_tensor_sh_encoder() diff --git a/tests/nerf/test_renderers.py b/tests/nerf/test_renderers.py index 8c2170509..4cbf6b147 100644 --- a/tests/nerf/test_renderers.py +++ b/tests/nerf/test_renderers.py @@ -23,8 +23,56 @@ def test_rgb_renderer(): out = rgb_renderer(rgb=rgb * 0, weights=weights) assert torch.max(out.rgb) == pytest.approx(0) - # TODO Tancik: Add more precise tests + +def test_sh_renderer(): + """Test SH volumetric rendering""" + + levels = 4 + num_samples = 10 + + sh = torch.ones((3, num_samples, 3 * levels**2)) + weights = torch.ones((3, num_samples)) + weights /= torch.sum(weights, axis=-1, keepdim=True) + directions = torch.zeros((3, num_samples, 3)) + directions[..., 0] = 1 + + sh_renderer = renderers.SHRenderer() + + out = sh_renderer(sh=sh, directions=directions, weights=weights) + assert torch.max(out.rgb) > 0.9 + + +def test_acc_renderer(): + """Test accumulation rendering""" + + num_samples = 10 + weights = torch.ones((3, num_samples)) + weights /= torch.sum(weights, axis=-1, keepdim=True) + + acc_renderer = renderers.AccumulationRenderer() + + out = acc_renderer(weights=weights) + assert torch.max(out.accumulation) > 0.9 + + +def test_depth_renderer(): + """Test depth rendering""" + + num_samples = 10 + weights = torch.ones((3, num_samples)) + weights /= torch.sum(weights, axis=-1, keepdim=True) + + ts = torch.linspace(0, 100, num_samples) + ts = torch.stack([ts, ts, ts], dim=0) + + depth_renderer = renderers.DepthRenderer() + + out = depth_renderer(weights=weights, ts=ts) + assert torch.min(out.depth) > 0 if __name__ == "__main__": test_rgb_renderer() + test_sh_renderer() + test_acc_renderer() + test_depth_renderer()