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": [
+ "![](\"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAJYAAABkCAIAAADrOV6nAAABA0lEQVR4nO3RsQkAIQDAwPcrG+d2bYewkMDdBIGMNfdH2f86gFsW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYZ2GehXkW5lmYdwDkygFkzaTKPwAAAABJRU5ErkJggg==\"/) |
"
+ ],
+ "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()