Add spherical harmoincs

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)

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)

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

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

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()

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()