From f183e03517c6307671efa817ea9aac91c18aa5ba Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Wed, 8 Jun 2022 11:43:16 -0700 Subject: [PATCH] Add gausian frustum calculations. --- pyrad/cameras/rays.py | 34 +++++++++++--- pyrad/utils/math.py | 96 ++++++++++++++++++++++++++++++++++++++ tests/cameras/test_rays.py | 74 +++++++++++++++++++++++++++++ 3 files changed, 197 insertions(+), 7 deletions(-) create mode 100644 tests/cameras/test_rays.py diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index 6df16a730..348ac68fa 100644 --- a/pyrad/cameras/rays.py +++ b/pyrad/cameras/rays.py @@ -7,6 +7,7 @@ from typing import Optional import torch from torchtyping import TensorType +from pyrad.utils.math import Gaussians, conical_frustum_to_gaussian from pyrad.utils.misc import is_not_none @@ -30,9 +31,29 @@ class Frustums: pixel_area: TensorType[..., 1] def get_positions(self) -> TensorType[..., 3]: - """Returns "center" position of frustum. Not weighted by mass.""" + """Calulates "center" position of frustum. Not weighted by mass. + + Returns: + TensorType[..., 3]: xyz positions. + """ return self.origins + self.directions * (self.frustum_starts + self.frustum_ends) / 2 + def get_gaussian_blob(self) -> Gaussians: + """Calculates of guassian approximation of conical frustum. + + Resturns: + Gaussians: Conical frustums approximated by gaussian distribution. + """ + # Cone radius is set such that the square pixel_area matches the cone area. + cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159 # r = sqrt(pixel_area / pi) + return conical_frustum_to_gaussian( + origins=self.origins, + directions=self.directions, + starts=self.frustum_starts, + ends=self.frustum_ends, + radius=cone_radius, + ) + @classmethod def get_mock_frustum(cls) -> "Frustums": """Helper function to generate a placeholder frustum. @@ -40,13 +61,12 @@ class Frustums: Returns: Frustums: A size 1 frustum with meaningless values. """ - device = cls.origins.device return Frustums( - origins=torch.ones((1, 3), device=device), - directions=torch.ones((1, 3), device=device), - frustum_starts=torch.ones((1, 1), device=device), - frustum_ends=torch.ones((1, 1), device=device) + 1, - pixel_area=torch.ones((1, 1), device=device), + origins=torch.ones((1, 3)), + directions=torch.ones((1, 3)), + frustum_starts=torch.ones((1, 1)), + frustum_ends=torch.ones((1, 1)), + pixel_area=torch.ones((1, 1)), ) def apply_masks(self, mask: TensorType) -> "Frustums": diff --git a/pyrad/utils/math.py b/pyrad/utils/math.py index 5bf1767f7..c28d5da62 100644 --- a/pyrad/utils/math.py +++ b/pyrad/utils/math.py @@ -1,5 +1,6 @@ """ Math Helper Functions """ +from dataclasses import dataclass import torch from torchtyping import TensorType @@ -67,3 +68,98 @@ def components_from_spherical_harmonics(levels: int, directions: TensorType[..., components[..., 24] = 0.4425326924449826 * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) return components + + +@dataclass +class Gaussians: + """Stores Gaussians + + Args: + mean (TensorType[..., "dim"]): Mean of multivariate Gaussian + cov (TensorType[..., "dim", "dim"]): Covariance of multivariate Gaussian. + """ + + mean: TensorType[..., "dim"] + cov: TensorType[..., "dim", "dim"] + + +def compute_3d_gaussian( + directions: TensorType[..., 3], + means: TensorType[..., 3], + dir_variance: TensorType[..., 1], + radius_variance: TensorType[..., 1], +) -> Gaussians: + """Compute guassian along ray. + + Args: + directions (TensorType[..., 3]): Axis of Gaussian. + means (TensorType[..., 3]): Mean of Gaussian. + dir_variance (TensorType[..., 3]): Variance along direction axis. + radius_variance (TensorType[..., 3]): Variance tangent to direction axis. + + Returns: + Gaussians: Oriented 3D gaussian. + """ + + dir_outer_product = directions[..., :, None] * directions[..., None, :] + eye = torch.eye(directions.shape[-1]) + dir_mag_sq = torch.clamp(torch.sum(directions**2, axis=-1, keepdim=True), min=1e-10) + null_outer_product = eye - directions[..., :, None] * (directions / dir_mag_sq)[..., None, :] + dir_cov_diag = dir_variance[..., None] * dir_outer_product[..., :, :] + radius_cov_diag = radius_variance[..., None] * null_outer_product[..., :, :] + cov = dir_cov_diag + radius_cov_diag + return Gaussians(mean=means, cov=cov) + + +def cylinder_to_gaussian( + origins: TensorType[..., 3], + directions: TensorType[..., 3], + starts: TensorType[..., 1], + ends: TensorType[..., 1], + radius: TensorType[..., 1], +) -> Gaussians: + """Approximates cylinders with a Gaussian distributions. + + Args: + origins (TensorType[..., 3]): Origins of cylinders. + directions (TensorType[..., 3]): Direction (axis) of cylinders. + starts (TensorType[..., 1]): Start of cylinders. + ends (TensorType[..., 1]): End of cylinders. + radius (TensorType[..., 1]): Radii of cylinders. + + Returns: + Gaussians: Approximation of cylinders + """ + means = origins + directions * ((starts + ends) / 2.0) + dir_variance = (ends - starts) ** 2 / 12 + radius_variance = radius**2 / 4.0 + return compute_3d_gaussian(directions, means, dir_variance, radius_variance) + + +def conical_frustum_to_gaussian( + origins: TensorType[..., 3], + directions: TensorType[..., 3], + starts: TensorType[..., 1], + ends: TensorType[..., 1], + radius: TensorType[..., 1], +) -> Gaussians: + """Approximates conical frustums with a Gaussian distributions. + + Uses stable parameterization described in mip-NeRF publication. + + Args: + origins (TensorType[..., 3]): Origins of cones. + directions (TensorType[..., 3]): Direction (axis) of frustums. + starts (TensorType[..., 1]): Start of conical frustums. + ends (TensorType[..., 1]): End of conical frustums. + radius (TensorType[..., 1]): Radii of cone a distance of 1 from the origin. + + Returns: + Gaussians: Approximation of conical frustums + """ + mu = (starts + ends) / 2.0 + hw = (ends - starts) / 2.0 + means = origins + directions * (mu + (2.0 * mu * hw**2.0) / (3.0 * mu**2.0 + hw**2.0)) + dir_variance = (hw**2) / 3 - (4 / 15) * ((hw**4 * (12 * mu**2 - hw**2)) / (3 * mu**2 + hw**2) ** 2) + radius_variance = radius**2 * ((mu**2) / 4 + (5 / 12) * hw**2 - 4 / 15 * (hw**4) / (3 * mu**2 + hw**2)) + return compute_3d_gaussian(directions, means, dir_variance, radius_variance) diff --git a/tests/cameras/test_rays.py b/tests/cameras/test_rays.py new file mode 100644 index 000000000..3b3d70c16 --- /dev/null +++ b/tests/cameras/test_rays.py @@ -0,0 +1,74 @@ +""" +Test the ray classes. +""" + +import pytest +import torch +from pyrad.cameras.rays import Frustums + + +def test_frustum_get_position(): + """Test position calculation""" + + origin = torch.Tensor([0, 1, 2])[None, ...] + direction = torch.Tensor([0, 1, 0])[None, ...] + frustum_start = torch.Tensor([2])[None, ...] + frustum_end = torch.Tensor([3])[None, ...] + + target_position = torch.Tensor([0, 3.5, 2])[None, ...] + + frustum = Frustums( + origins=origin, + directions=direction, + frustum_starts=frustum_start, + frustum_ends=frustum_end, + pixel_area=torch.ones((1, 1)), + ) + + positions = frustum.get_positions() + assert positions == pytest.approx(target_position, abs=1e-6) + + +def test_frustum_get_gaussian_blob(): + """Test gaussian blob calculation""" + + frustum = Frustums( + origins=torch.ones((5, 3)), + directions=torch.ones((5, 3)), + frustum_starts=torch.ones((5, 1)), + frustum_ends=torch.ones((5, 1)), + pixel_area=torch.ones((5, 1)), + ) + + gaussian_blob = frustum.get_gaussian_blob() + assert gaussian_blob.mean.shape == (5, 3) + assert gaussian_blob.cov.shape == (5, 3, 3) + + +def test_frustum_apply_masks(): + """Test masking frustum""" + frustum = Frustums( + origins=torch.ones((5, 3)), + directions=torch.ones((5, 3)), + frustum_starts=torch.ones((5, 1)), + frustum_ends=torch.ones((5, 1)), + pixel_area=torch.ones((5, 1)), + ) + + mask = torch.tensor([False, True, False, True, True], dtype=torch.bool) + frustum = frustum.apply_masks(mask) + + assert frustum.origins.shape == (3, 3) + assert frustum.directions.shape == (3, 3) + assert frustum.frustum_starts.shape == (3, 1) + assert frustum.frustum_ends.shape == (3, 1) + assert frustum.pixel_area.shape == (3, 1) + + +def test_get_mock_frustum(): + """Test creation of mock frustum""" + Frustums.get_mock_frustum() + + +if __name__ == "__main__": + test_frustum_get_gaussian_blob()