286 lines
9.7 KiB
Python
286 lines
9.7 KiB
Python
# 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 numpy as np
|
|
|
|
import torch
|
|
|
|
from pytorch3d.implicitron.models.renderer.base import (
|
|
approximate_conical_frustum_as_gaussians,
|
|
compute_3d_diagonal_covariance_gaussian,
|
|
conical_frustum_to_gaussian,
|
|
ImplicitronRayBundle,
|
|
)
|
|
from pytorch3d.implicitron.models.renderer.ray_sampler import AbstractMaskRaySampler
|
|
|
|
from tests.common_testing import TestCaseMixin
|
|
|
|
|
|
class TestRendererBase(TestCaseMixin, unittest.TestCase):
|
|
def test_implicitron_from_bins(self) -> None:
|
|
bins = torch.randn(2, 3, 4, 5)
|
|
ray_bundle = ImplicitronRayBundle(
|
|
origins=None,
|
|
directions=None,
|
|
lengths=None,
|
|
xys=None,
|
|
bins=bins,
|
|
)
|
|
self.assertClose(ray_bundle.lengths, 0.5 * (bins[..., 1:] + bins[..., :-1]))
|
|
self.assertClose(ray_bundle.bins, bins)
|
|
|
|
def test_implicitron_raise_value_error_bins_is_set_and_try_to_set_lengths(
|
|
self,
|
|
) -> None:
|
|
ray_bundle = ImplicitronRayBundle(
|
|
origins=torch.rand(2, 3, 4, 3),
|
|
directions=torch.rand(2, 3, 4, 3),
|
|
lengths=None,
|
|
xys=torch.rand(2, 3, 4, 2),
|
|
bins=torch.rand(2, 3, 4, 14),
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"If the bins attribute is not None you cannot set the lengths attribute.",
|
|
):
|
|
ray_bundle.lengths = torch.empty(2)
|
|
|
|
def test_implicitron_raise_value_error_if_bins_dim_equal_1(self) -> None:
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The last dim of bins must be at least superior or equal to 2."
|
|
):
|
|
ImplicitronRayBundle(
|
|
origins=torch.rand(2, 3, 4, 3),
|
|
directions=torch.rand(2, 3, 4, 3),
|
|
lengths=None,
|
|
xys=torch.rand(2, 3, 4, 2),
|
|
bins=torch.rand(2, 3, 4, 1),
|
|
)
|
|
|
|
def test_implicitron_raise_value_error_if_neither_bins_or_lengths_provided(
|
|
self,
|
|
) -> None:
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"Please set either bins or lengths to initialize an ImplicitronRayBundle.",
|
|
):
|
|
ImplicitronRayBundle(
|
|
origins=torch.rand(2, 3, 4, 3),
|
|
directions=torch.rand(2, 3, 4, 3),
|
|
lengths=None,
|
|
xys=torch.rand(2, 3, 4, 2),
|
|
bins=None,
|
|
)
|
|
|
|
def test_conical_frustum_to_gaussian(self) -> None:
|
|
origins = torch.zeros(3, 3, 3)
|
|
directions = torch.tensor(
|
|
[
|
|
[[0, 0, 0], [1, 0, 0], [3, 0, 0]],
|
|
[[0, 0.25, 0], [1, 0.25, 0], [3, 0.25, 0]],
|
|
[[0, 1, 0], [1, 1, 0], [3, 1, 0]],
|
|
]
|
|
)
|
|
bins = torch.tensor(
|
|
[
|
|
[[0.5, 1.5], [0.3, 0.7], [0.3, 0.7]],
|
|
[[0.5, 1.5], [0.3, 0.7], [0.3, 0.7]],
|
|
[[0.5, 1.5], [0.3, 0.7], [0.3, 0.7]],
|
|
]
|
|
)
|
|
# see test_compute_pixel_radii_from_ray_direction
|
|
radii = torch.tensor(
|
|
[
|
|
[1.25, 2.25, 2.25],
|
|
[1.75, 2.75, 2.75],
|
|
[1.75, 2.75, 2.75],
|
|
]
|
|
)
|
|
radii = radii[..., None] / 12**0.5
|
|
|
|
# The expected mean and diagonal covariance have been computed
|
|
# by hand from the official code of MipNerf.
|
|
# https://github.com/google/mipnerf/blob/84c969e0a623edd183b75693aed72a7e7c22902d/internal/mip.py#L125
|
|
# mean, cov_diag = cast_rays(length, origins, directions, radii, 'cone', diag=True)
|
|
|
|
expected_mean = torch.tensor(
|
|
[
|
|
[
|
|
[[0.0, 0.0, 0.0]],
|
|
[[0.5506329, 0.0, 0.0]],
|
|
[[1.6518986, 0.0, 0.0]],
|
|
],
|
|
[
|
|
[[0.0, 0.28846154, 0.0]],
|
|
[[0.5506329, 0.13765822, 0.0]],
|
|
[[1.6518986, 0.13765822, 0.0]],
|
|
],
|
|
[
|
|
[[0.0, 1.1538461, 0.0]],
|
|
[[0.5506329, 0.5506329, 0.0]],
|
|
[[1.6518986, 0.5506329, 0.0]],
|
|
],
|
|
]
|
|
)
|
|
expected_diag_cov = torch.tensor(
|
|
[
|
|
[
|
|
[[0.04544772, 0.04544772, 0.04544772]],
|
|
[[0.01130973, 0.03317059, 0.03317059]],
|
|
[[0.10178753, 0.03317059, 0.03317059]],
|
|
],
|
|
[
|
|
[[0.08907752, 0.00404956, 0.08907752]],
|
|
[[0.0142245, 0.04734321, 0.04955113]],
|
|
[[0.10212927, 0.04991625, 0.04955113]],
|
|
],
|
|
[
|
|
[[0.08907752, 0.0647929, 0.08907752]],
|
|
[[0.03608529, 0.03608529, 0.04955113]],
|
|
[[0.10674264, 0.05590574, 0.04955113]],
|
|
],
|
|
]
|
|
)
|
|
|
|
ray = ImplicitronRayBundle(
|
|
origins=origins,
|
|
directions=directions,
|
|
bins=bins,
|
|
lengths=None,
|
|
pixel_radii_2d=radii,
|
|
xys=None,
|
|
)
|
|
mean, diag_cov = conical_frustum_to_gaussian(ray)
|
|
|
|
self.assertClose(mean, expected_mean)
|
|
self.assertClose(diag_cov, expected_diag_cov)
|
|
|
|
def test_scale_conical_frustum_to_gaussian(self) -> None:
|
|
origins = torch.zeros(2, 2, 3)
|
|
directions = torch.Tensor(
|
|
[
|
|
[[0, 1, 0], [0, 0, 1]],
|
|
[[0, 1, 0], [0, 0, 1]],
|
|
]
|
|
)
|
|
bins = torch.Tensor(
|
|
[
|
|
[[0.5, 1.5], [0.3, 0.7]],
|
|
[[0.5, 1.5], [0.3, 0.7]],
|
|
]
|
|
)
|
|
radii = torch.ones(2, 2, 1)
|
|
|
|
ray = ImplicitronRayBundle(
|
|
origins=origins,
|
|
directions=directions,
|
|
bins=bins,
|
|
pixel_radii_2d=radii,
|
|
lengths=None,
|
|
xys=None,
|
|
)
|
|
|
|
mean, diag_cov = conical_frustum_to_gaussian(ray)
|
|
|
|
scaling_factor = 2.5
|
|
ray = ImplicitronRayBundle(
|
|
origins=origins,
|
|
directions=directions,
|
|
bins=bins * scaling_factor,
|
|
pixel_radii_2d=radii,
|
|
lengths=None,
|
|
xys=None,
|
|
)
|
|
mean_scaled, diag_cov_scaled = conical_frustum_to_gaussian(ray)
|
|
np.testing.assert_allclose(mean * scaling_factor, mean_scaled)
|
|
np.testing.assert_allclose(
|
|
diag_cov * scaling_factor**2, diag_cov_scaled, atol=1e-6
|
|
)
|
|
|
|
def test_approximate_conical_frustum_as_gaussian(self) -> None:
|
|
"""Ensure that the computation modularity in our function is well done."""
|
|
bins = torch.Tensor([[0.5, 1.5], [0.3, 0.7]])
|
|
radii = torch.Tensor([[1.0], [1.0]])
|
|
t_mean, t_var, r_var = approximate_conical_frustum_as_gaussians(bins, radii)
|
|
|
|
self.assertEqual(t_mean.shape, (2, 1))
|
|
self.assertEqual(t_var.shape, (2, 1))
|
|
self.assertEqual(r_var.shape, (2, 1))
|
|
|
|
mu = np.array([[1.0], [0.5]])
|
|
delta = np.array([[0.5], [0.2]])
|
|
|
|
np.testing.assert_allclose(
|
|
mu + (2 * mu * delta**2) / (3 * mu**2 + delta**2), t_mean.numpy()
|
|
)
|
|
np.testing.assert_allclose(
|
|
(delta**2) / 3
|
|
- (4 / 15)
|
|
* ((delta**4 * (12 * mu**2 - delta**2)) / (3 * mu**2 + delta**2) ** 2),
|
|
t_var.numpy(),
|
|
)
|
|
np.testing.assert_allclose(
|
|
radii**2
|
|
* (
|
|
(mu**2) / 4
|
|
+ (5 / 12) * delta**2
|
|
- 4 / 15 * (delta**4) / (3 * mu**2 + delta**2)
|
|
),
|
|
r_var.numpy(),
|
|
)
|
|
|
|
def test_compute_3d_diagonal_covariance_gaussian(self) -> None:
|
|
ray_directions = torch.Tensor([[0, 0, 1]])
|
|
t_var = torch.Tensor([0.5, 0.5, 1])
|
|
r_var = torch.Tensor([0.6, 0.3, 0.4])
|
|
expected_diag_cov = np.array(
|
|
[
|
|
[
|
|
# t_cov_diag + xy_cov_diag
|
|
[0.0 + 0.6, 0.0 + 0.6, 0.5 + 0.0],
|
|
[0.0 + 0.3, 0.0 + 0.3, 0.5 + 0.0],
|
|
[0.0 + 0.4, 0.0 + 0.4, 1.0 + 0.0],
|
|
]
|
|
]
|
|
)
|
|
diag_cov = compute_3d_diagonal_covariance_gaussian(ray_directions, t_var, r_var)
|
|
np.testing.assert_allclose(diag_cov.numpy(), expected_diag_cov)
|
|
|
|
def test_conical_frustum_to_gaussian_raise_valueerror(self) -> None:
|
|
lengths = torch.linspace(0, 1, steps=6)
|
|
directions = torch.tensor([0, 0, 1])
|
|
origins = torch.tensor([1, 1, 1])
|
|
ray = ImplicitronRayBundle(
|
|
origins=origins, directions=directions, lengths=lengths, xys=None
|
|
)
|
|
|
|
expected_error_message = (
|
|
"RayBundle pixel_radii_2d or bins have not been provided."
|
|
" Look at pytorch3d.renderer.implicit.renderer.ray_sampler::"
|
|
"AbstractMaskRaySampler to see how to compute them. Have you forgot to set"
|
|
"`cast_ray_bundle_as_cone` to True?"
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, expected_error_message):
|
|
_ = conical_frustum_to_gaussian(ray)
|
|
|
|
# Ensure message is coherent with AbstractMaskRaySampler
|
|
class FakeRaySampler(AbstractMaskRaySampler):
|
|
def _get_min_max_depth_bounds(self, *args):
|
|
return None
|
|
|
|
message_assertion = (
|
|
"If cast_ray_bundle_as_cone has been removed please update the doc"
|
|
"conical_frustum_to_gaussian"
|
|
)
|
|
self.assertIsNotNone(
|
|
getattr(FakeRaySampler(), "cast_ray_bundle_as_cone", None),
|
|
message_assertion,
|
|
)
|