Distortion loss (#58)

* add distortion loss

* update init

* bump version

nerfacc/__init__.py

@@ -7,6 +7,7 @@
 from .contraction import ContractionType, contract, contract_inv
 from .grid import Grid, OccupancyGrid, query_grid
 from .intersection import ray_aabb_intersect
+from .losses import distortion as loss_distortion
 from .pack import pack_data, unpack_data, unpack_info
 from .ray_marching import ray_marching
 from .version import __version__
@@ -48,4 +49,5 @@ def unpack_to_ray_indices(*args, **kwargs):
     "unpack_data",
     "unpack_info",
     "ray_resampling",
+    "loss_distortion",
 ]

nerfacc/losses.py

@@ -0,0 +1,32 @@
+from torch import Tensor
+
+from .pack import unpack_data
+
+
+def distortion(
+    packed_info: Tensor, weights: Tensor, t_starts: Tensor, t_ends: Tensor
+) -> Tensor:
+    """Distortion loss from Mip-NeRF 360 paper, Equ. 15.
+
+    Args:
+        packed_info: Packed info for the samples. (n_rays, 2)
+        weights: Weights for the samples. (all_samples,)
+        t_starts: Per-sample start distance. Tensor with shape (all_samples, 1).
+        t_ends: Per-sample end distance. Tensor with shape (all_samples, 1).
+
+    Returns:
+        Distortion loss. (n_rays,)
+    """
+    # （all_samples, 1) -> (n_rays, n_samples)
+    w = unpack_data(packed_info, weights[..., None]).squeeze(-1)
+    t1 = unpack_data(packed_info, t_starts).squeeze(-1)
+    t2 = unpack_data(packed_info, t_ends).squeeze(-1)
+
+    interval = t2 - t1
+    tmid = (t1 + t2) / 2
+
+    loss_uni = (1 / 3) * (interval * w.pow(2)).sum(-1)
+    ww = w.unsqueeze(-1) * w.unsqueeze(-2)
+    mm = (tmid.unsqueeze(-1) - tmid.unsqueeze(-2)).abs()
+    loss_bi = (ww * mm).sum((-1, -2))
+    return loss_uni + loss_bi

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nerfacc"
-version = "0.1.7"
+version = "0.1.8"
 description = "A General NeRF Acceleration Toolbox."
 readme = "README.md"
 authors = [{name = "Ruilong", email = "[email protected]"}]

tests/test_loss.py

@@ -0,0 +1,31 @@
+import pytest
+import torch
+
+from nerfacc import ray_marching
+from nerfacc.losses import distortion
+
+device = "cuda:0"
+batch_size = 32
+eps = 1e-6
+
+
+@pytest.mark.skipif(not torch.cuda.is_available, reason="No CUDA device")
+def test_distortion():
+    rays_o = torch.rand((batch_size, 3), device=device)
+    rays_d = torch.randn((batch_size, 3), device=device)
+    rays_d = rays_d / rays_d.norm(dim=-1, keepdim=True)
+
+    packed_info, t_starts, t_ends = ray_marching(
+        rays_o,
+        rays_d,
+        near_plane=0.1,
+        far_plane=1.0,
+        render_step_size=1e-3,
+    )
+    weights = torch.rand((t_starts.shape[0],), device=device)
+    loss = distortion(packed_info, weights, t_starts, t_ends)
+    assert loss.shape == (batch_size,)
+
+
+if __name__ == "__main__":
+    test_distortion()