Add integrated position encoding based on MIPNerf implementation.

Summary: Add a new implicit module Integral Position Encoding based on [MIP-NeRF](https://arxiv.org/abs/2103.13415). Reviewed By: shapovalov Differential Revision: D46352730 fbshipit-source-id: c6a56134c975d80052b3a11f5e92fd7d95cbff1e
2025-12-20 14:20:38 +08:00 · 2023-07-06 02:20:53 -07:00
parent 29b8ebd802
commit ccf860f1db
6 changed files with 283 additions and 29 deletions
--- a/tests/implicitron/test_implicit_function_neural_radiance_field.py
+++ b/tests/implicitron/test_implicit_function_neural_radiance_field.py
@@ -0,0 +1,66 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from pytorch3d.implicitron.models.implicit_function.base import ImplicitronRayBundle
+from pytorch3d.implicitron.models.implicit_function.neural_radiance_field import (
+    NeuralRadianceFieldImplicitFunction,
+)
+
+
+class TestNeuralRadianceFieldImplicitFunction(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+
+    def test_forward_with_integrated_positionial_embedding(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            bins=torch.randn(*shape, 6 + 1),
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(
+            n_hidden_neurons_dir=32, use_integrated_positional_encoding=True
+        )
+        raw_densities, ray_colors, _ = model(ray_bundle=ray_bundle)
+
+        self.assertEqual(raw_densities.shape, (*shape, ray_bundle.lengths.shape[-1], 1))
+        self.assertEqual(ray_colors.shape, (*shape, ray_bundle.lengths.shape[-1], 3))
+
+    def test_forward_with_integrated_positionial_embedding_raise_exception(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            bins=None,
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(
+            n_hidden_neurons_dir=32, use_integrated_positional_encoding=True
+        )
+        with self.assertRaises(ValueError):
+            _ = model(ray_bundle=ray_bundle)
+
+    def test_forward(self):
+        shape = [2, 4, 4]
+        ray_bundle = ImplicitronRayBundle(
+            origins=torch.randn(*shape, 3),
+            directions=torch.randn(*shape, 3),
+            lengths=torch.randn(*shape, 6),
+            pixel_radii_2d=torch.randn(*shape, 1),
+            xys=None,
+        )
+        model = NeuralRadianceFieldImplicitFunction(n_hidden_neurons_dir=32)
+        raw_densities, ray_colors, _ = model(ray_bundle=ray_bundle)
+        self.assertEqual(raw_densities.shape, (*shape, 6, 1))
+        self.assertEqual(ray_colors.shape, (*shape, 6, 3))
--- a/tests/test_harmonic_embedding.py
+++ b/tests/test_harmonic_embedding.py
@@ -8,6 +8,7 @@ import unittest

 import torch
 from pytorch3d.renderer.implicit import HarmonicEmbedding
+from torch.distributions import MultivariateNormal

 from .common_testing import TestCaseMixin

@@ -36,16 +37,117 @@ class TestHarmonicEmbedding(TestCaseMixin, unittest.TestCase):
        self.assertClose(embed_fun_lin._frequencies, torch.FloatTensor((1.0, 2.5, 4.0)))

    def test_correct_embed_out(self):
-        embed_fun = HarmonicEmbedding(n_harmonic_functions=2, append_input=False)
+        n_harmonic_functions = 2
        x = torch.randn((1, 5))
-        D = 5 * 4
+        D = 5 * n_harmonic_functions * 2  # sin + cos
+
+        embed_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions, append_input=False
+        )
        embed_out = embed_fun(x)
+
        self.assertEqual(embed_out.shape, (1, D))
        # Sum the squares of the respective frequencies
+        # cos^2(x) + sin^2(x) = 1
        sum_squares = embed_out[0, : D // 2] ** 2 + embed_out[0, D // 2 :] ** 2
        self.assertClose(sum_squares, torch.ones((D // 2)))
-        embed_fun = HarmonicEmbedding(n_harmonic_functions=2, append_input=True)
-        embed_out = embed_fun(x)
-        self.assertClose(embed_out.shape, torch.tensor((1, 5 * 5)))
+
+        # Test append input
+        embed_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions, append_input=True
+        )
+        embed_out_appended_input = embed_fun(x)
+        self.assertClose(
+            embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
+        )
        # Last plane in output is the input
-        self.assertClose(embed_out[..., -5:], x)
+        self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
+        self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)
+
+    def test_correct_embed_out_with_diag_cov(self):
+        n_harmonic_functions = 2
+        x = torch.randn((1, 3))
+        diag_cov = torch.randn((1, 3))
+        D = 3 * n_harmonic_functions * 2  # sin + cos
+
+        embed_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions, append_input=False
+        )
+        embed_out = embed_fun(x, diag_cov=diag_cov)
+
+        self.assertEqual(embed_out.shape, (1, D))
+
+        # Compute the scaling factor introduce in MipNerf
+        scale_factor = (
+            -0.5 * diag_cov[..., None] * torch.pow(embed_fun._frequencies[None, :], 2)
+        )
+        scale_factor = torch.exp(scale_factor).reshape(1, -1).tile((1, 2))
+        # If we remove this scaling factor, we should go back to the
+        # classical harmonic embedding:
+        # Sum the squares of the respective frequencies
+        # cos^2(x) + sin^2(x) = 1
+        embed_out_without_cov = embed_out / scale_factor
+        sum_squares = (
+            embed_out_without_cov[0, : D // 2] ** 2
+            + embed_out_without_cov[0, D // 2 :] ** 2
+        )
+        self.assertClose(sum_squares, torch.ones((D // 2)))
+
+        # Test append input
+        embed_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions, append_input=True
+        )
+        embed_out_appended_input = embed_fun(x, diag_cov=diag_cov)
+        self.assertClose(
+            embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
+        )
+        # Last plane in output is the input
+        self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
+        self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)
+
+    def test_correct_behavior_between_ipe_and_its_estimation_from_harmonic_embedding(
+        self,
+    ):
+        """
+        Check that the HarmonicEmbedding with integrated_position_encoding (IPE) set to
+        True is coherent with the HarmonicEmbedding.
+
+        What is the idea behind this test?
+
+        We wish to produce an IPE that is the expectation
+        of our lifted multivariate gaussian, modulated by the sine and cosine of
+        the coordinates. These expectation has a closed-form
+        (see equations 11, 12, 13, 14 of [1]).
+
+        We sample N elements from the multivariate gaussian defined by its mean and covariance
+        and compute the HarmonicEmbedding. The expected value of those embeddings should be
+        equal to our IPE.
+
+        Inspired from:
+        https://github.com/google/mipnerf/blob/84c969e0a623edd183b75693aed72a7e7c22902d/internal/mip_test.py#L359
+
+        References:
+            [1] `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
+        """
+        num_dims = 3
+        n_harmonic_functions = 6
+        mean = torch.randn(num_dims)
+        diag_cov = torch.rand(num_dims)
+
+        he_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions, logspace=True, append_input=False
+        )
+        ipe_fun = HarmonicEmbedding(
+            n_harmonic_functions=n_harmonic_functions,
+            append_input=False,
+        )
+
+        embedding_ipe = ipe_fun(mean, diag_cov=diag_cov)
+
+        rand_mvn = MultivariateNormal(mean, torch.eye(num_dims) * diag_cov)
+
+        # Providing a large enough number of samples
+        # we should obtain an estimation close to our IPE
+        num_samples = 100000
+        embedding_he = he_fun(rand_mvn.sample_n(num_samples))
+        self.assertClose(embedding_he.mean(0), embedding_ipe, rtol=1e-2, atol=1e-2)